0x01 データ読み込み
1.1 加速経路
分散トレーニングでは、トレーニングを高速化するために、3 つのレベルの作業に対処する必要があります。
- データ読み込みレベル
- マルチマシン通信レベル
- コードレベル
データ レベルでは、マルチプロセスの並列読み込みを使用してデータの前処理プロセスを高速化することができ、GPU 機能を使用して高速化することもできます.たとえば、Nvidia DALI はデータの前処理を GPU に配置することで CPU ボトルネックの問題を解決します.
マルチマシン通信レベルでは、NCCL、OpenMPI、Gloo など、さまざまな集合通信ライブラリが利用できます。
コード レベルでは、フレームワークによって提供される分散 API を使用するか、Horovod を使用してスタンドアロン コードを変更し、分散タスクをサポートすることができます。
次に、データレベルを高速化する方法を見てみましょう。
1.2 並列処理
AIフレームワークのデータ処理は、主に以下のような並列処理です。
- データのロード/処理には CPU が使用されます。
- トレーニングは GPU を使用します。
理想的には、反復トレーニングの各ラウンドの前に、CPU がロードされ、トレーニング データの準備が整っているため、トレーニングをシームレスに反復し続けることができます。
ただし、GPUの演算能力は毎年2倍になり、CPUの向上速度はGPUに大きく遅れをとっているため、CPUが足を引っ張っています。これは、CPU の計算能力が不足しているという問題だけでなく、村のストレージでのデータの読み込み速度が不十分であるという問題でもあります。
したがって、機械学習では、データの読み込みと早期の前処理に対する要件がますます高くなり、データの次の反復の準備は GPU の計算時間内に完了する必要があり、GPU はトレーニング データを待機しているため、アイドル状態になることはできません。
1.3 パイプライン
機械学習トレーニングの場合、データの読み込みは次の 3 つのステップに分けることができます。
- ディスクまたは分散ストレージからホスト (CPU) にデータをロードします。
- ホストのページング可能なメモリからホストの固定メモリにデータを転送します。
- ホスト固定メモリからホスト GPU にデータを転送します。
したがって、一般的な深層学習フレームワークは、ロード ステップの特性と異種ハードウェアの特性に従ってパイプライン処理を実行し、それによってデータ処理プロセスのスループットを向上させます。
パイプラインには通常複数のオペレーターが含まれ、各オペレーターはバッファーを形成するデータ キューで構成され、上流のオペレーターが処理を完了すると、処理のために下流のオペレーターに渡されます。このように、各オペレーターのタスクは互いに独立しています. オペレーターは、細粒度のマルチスレッド/マルチプロセスを使用して並列に高速化できます. 各オペレーターは、処理速度とメモリを独立して制御して、処理に適応できます.さまざまなネットワークの速度要件。
オペレーターの内部データ キューが空でない場合、モデルは継続的にデータを取得し、トレーニング データの待機によるボトルネックは発生しません。
シリアル処理ロジックは次のとおりです。
+------+ +-----------+ +---------------------------+
| | | | | |
| Data +----------> | Load Data +---------> | Transfer to Pinned Memory |
| | | | | |
+------+ +-----------+ +---------------------------+
以下は、並列パイプライン ロジックです。
+------------+
+--------+ | |
| | | Process 1 |
| Data 1 +--------> | +------+
| | | Load Data | |
+--------+ | | |
+------------+ |
|
|
|
+------------+ | +-----------------------------------+
+--------+ | | | | |
| | | Process 2 | +------> | Pin-memory process |
| Data 2 +--------> | | | |
| | | Load Data +-------------> | |
+--------+ | | | Transfer to Pinned Memory |
+------------+ +-----> | |
| | |
| +-----------------------------------+
|
+--------+ +------------+ |
| | | | |
| Data 3 +--------> | Process 3 +-------+
| | | |
+--------+ | Load Data |
| |
+------------+
1.4 GPU
この記事ではこれまでのところ、CPU 側でのデータ転送の問題、つまりディスクからのデータのロード、ページング可能なメモリから固定メモリへのデータ転送について説明しました。
ただし、固定メモリから GPU へのデータ転送 ( ) は、tensor.cuda()CUDA ストリームを使用してパイプライン化することもできます。
さらに、ディープ ラーニング アプリケーションには、読み込み、デコード、クロッピング、サイズ変更、およびその他の多くの拡張機能を含む、複雑なマルチステージ データ処理パイプラインが必要です。現在 CPU で実行されているこれらのデータ処理パイプラインはボトルネックとなり、トレーニングと推論のパフォーマンスとスケーラビリティを制限しています。
Nvidia DALI は、データの前処理を GPU 処理に入れることで CPU ボトルネックの問題を解決し、ユーザーはモデルの特性に応じて GPU ベースのパイプラインまたは CPU ベースのパイプラインを構築できます。
次に、PyTorch のデータ読み込みについて、主に配布の観点から紹介します。
0x02 PyTorch 分散ローディング
2.1 DDP
Pytorch は、データ分散トレーニングのためのさまざまなオプションを提供します。アプリケーションが単純なものから複雑なものへ、プロトタイプから実稼働へと進むにつれて、一般的な開発の軌跡は次のようになります。
- データとモデルを単一の GPU に入れ、単一のデバイスでトレーニングできる場合、現時点ではトレーニング速度について心配する必要はありません。
- サーバーに複数の GPU があり、最小限のコード変更でトレーニングを高速化する場合は、複数の GPU DataParallel を備えた単一のマシンを使用します。
- トレーニングをさらに高速化し、開始するコードを少し書いても構わない場合は、複数の GPU DistributedDataParallel を備えた単一のマシンを使用します。
- アプリケーションがマシンの境界を越えてスケーリングする場合は、複数マシンの DistributedDataParallel と起動スクリプトを使用します。
- エラーが予想される場合 (OOM など)、またはトレーニング中にリソースを動的にアタッチおよびデタッチできる場合は、torchelastic を使用して分散トレーニングを有効にします。
この記事の最も重要な部分は DDP です。分散データ並列トレーニング (DDP) は、広く使用されている単一プログラムのマルチデータ トレーニング方法です。DDP を使用すると、モデルが各プロセスに複製され、モデルの各コピーにデータ サンプルの異なるサブセットが供給されます。DDP は勾配通信を処理してモデル レプリカの同期を維持し、それを勾配計算とオーバーラップしてトレーニングを高速化します。
2.2 分散負荷
まず、分散ローディングの一般的な構造を見ていきます。
サンプル コードを見ると、DataSet、DistributedSampler、DataLoader が主に使用されていることがわかります。
sampler = DistributedSampler(dataset) if is_distributed else None
loader = DataLoader(dataset, shuffle=(sampler is None), sampler=sampler)
for epoch in range(start_epoch, n_epochs):
if is_distributed:
sampler.set_epoch(epoch)
train(loader)
これら 3 つの概念の論理的な関係は次のとおりです。
- データセット: 名前からわかるように、データセットを意味します。元のトレーニング データをカプセル化し、それを Python で認識可能なデータ構造にカプセル化する責任を負います. Dataset の派生クラスは、単一のデータを取得するためのインターフェイスを提供する必要があります.
- サンプラー: 名前からわかるように、サンプリング方法またはサンプリング戦略を担当するサンプラーであり、特定の抽出/サンプリング戦略を実装して、DataLoade が使用するデータセットからデータ インデックスを取得します。サンプラーは、戦闘の場所を決定する司令官だと考えてください。
- DataLoader : インデックスに従ってデータセットからデータをロードする責任があります。Map スタイルおよび Iterable スタイルのデータセットをサポートし、シングルプロセス/マルチプロセスの読み込みをサポートします。ローダーは特定の戦闘機であり、サンプラーの命令に従って戦う責任があります。
具体的には、以下の図に示すように、簡単に説明します。
- DataSet は、データ セット番号を DistributedSampler に送信します。
- Sampler は、特定のルールに従ってデータ インデックスを Loader に送信します。
- ローダーは、インデックスに従ってデータをロードします。
- ローダーは、トレーニングのためにデータをモデルに送信します。
+------------------------+ +-----------+
|DistributedSampler | |DataLoader |
| | 2 indices | |
| Some strategy +-------------------> | |
| | | |
|-------------+----------| | |
^ | | 4 data +-------+
| | -------------->+ train |
1 | length | | +-------+
| | |
+-------------+----------+ | |
|DataSet | | |
| +---------+ | 3 Load | |
| | Data +-------------------------> | |
| +---------+ | | |
| | | |
+------------------------+ +-----------+
データセットは分散トレーニングの焦点ではないため、この記事では主に Sampler を次に分析します。
Sampler の焦点は、各ワーカーがデータ セット内の独自の部分のみをロードできるようにし、ワーカー間でデータ セットの直交割り当てを実現する方法です。
0x03 分散サンプラー
データの並列および分散トレーニングの場合、DistributedSampler はそのデータ サンプリングタスク。
DistributedSampler は Sampler の派生クラスです。DistributedDataParallel が DistributedSampler を使用する場合、各並列プロセスは DistributedSampler インスタンスを取得し、この DistributedSampler インスタンスは DataLoader に命令を送信して、DataLoader が特定のデータをロードするようにします。
DistributedSampler の読み込み戦略は、読み込まれたデータ セットのサブセットのみを提供する責任があり、DistributedSampler によって提供されるサブセットは重複したり重複したりしません。
3.1 初期化
__init__初期化コードは、主に、データセット、ランク (グローバル GPU シリアル番号)、num_replicas コピー数など、ワーカー ノードのさまざまな情報を設定します。そして、すべてのサンプルの total_size を計算します。
いくつかのパラメータは次のとおりです。
- dataset: サンプリングされたデータセットです。
- num_replicas : 分散トレーニングに参加しているプロセスの数. 設定されていない場合, world_size はプロセスの数としてグループから取得されます.
- rank : 現在のプロセスのシリアル番号。設定されていない場合は、グループから取得されます。
- shuffle : サンプリングでインデックスをシャッフルする必要があるかどうか。
- シード: スクランブルが必要な場合は、ランダム シードを設定します。
- drop_last : データを均等に分割できない場合、割り当てられない末尾のデータを破棄するかどうか。
- エポック: すべてのエポックでデータセットがシャッフルされます。シャッフル後にデータセットの一貫性を維持するにはどうすればよいですか? それはエポックを通して行われます。
具体的なコードは、例外処理を省略して次のとおりです。
class DistributedSampler(Sampler[T_co]):
def __init__(self, dataset: Dataset, num_replicas: Optional[int] = None,
rank: Optional[int] = None, shuffle: bool = True,
seed: int = 0, drop_last: bool = False) -> None:
self.dataset = dataset
self.num_replicas = num_replicas
self.rank = rank
self.epoch = 0
self.drop_last = drop_last
# If the dataset length is evenly divisible by # of replicas, then there
# is no need to drop any data, since the dataset will be split equally.
if self.drop_last and len(self.dataset) % self.num_replicas != 0: # type: ignore[arg-type]
# Split to nearest available length that is evenly divisible.
# This is to ensure each rank receives the same amount of data when
# using this Sampler.
self.num_samples = math.ceil(
# `type:ignore` is required because Dataset cannot provide a default __len__
# see NOTE in pytorch/torch/utils/data/sampler.py
(len(self.dataset) - self.num_replicas) / self.num_replicas # type: ignore[arg-type]
)
else:
self.num_samples = math.ceil(len(self.dataset) / self.num_replicas) # type: ignore[arg-type]
self.total_size = self.num_samples * self.num_replicas
self.shuffle = shuffle
self.seed = seed
3.2 反復法
DistributedSampler は (ループに似た) イテレーターとして実装されるため、python 抽象クラスのマジック メソッドを使用します。
__len__(self):len()関数、通常、イテレータ内の要素の数を返します。__iter__(self): コンテナー内の要素を反復処理するときの動作は、実際には反復子 (通常は反復子自体) を返し、各反復の結果は次の反復の初期値として使用されます。
__iter__コードの技術的な詳細は次のとおりです。
indices = indices[self.rank:self.total_size:self.num_replicas]
たとえば、リストに二重引用符がある場合、list[start:end:step]その意味は次のとおりです。
- 開始：開始位置
- end: 終了位置
- step: ステップサイズ
次のような例を見てみましょう。
a = [1,2,3,4,5,6,7,8,9,10,11,12,13,14,15]
print(a[0:15:3])
print(a[1:15:3])
print(a[2:15:3])
得る：
[1, 4, 7, 10, 13]
[2, 5, 8, 11, 14]
[3, 6, 9, 12, 15]
indices[self.rank:self.total_size:self.num_replicas]のnum_replicas は実際にはランクの総数であるため、ここでは各ワーカーは自身のランクに対応するデータ シーケンス番号の部分を厳密に返します。
要約すると、DistributedSampler の分散方法は次のとおりです。データの各連続num_replicasセグメントを1 つずつ分割num_replicasしてプロセスに分散し、各ワーカーのランクを通じてデータを取得して、非重複および非重複の目的を達成します。非交差だけでなく、各プロセスで取得されたデータは不連続であることに注意してください。
具体的なコードは次のとおりです。
class DistributedSampler(Sampler[T_co]):
def __iter__(self) -> Iterator[T_co]:
if self.shuffle: # 如果需要shuffle，则会基于epoch和seed进行处理
# deterministically shuffle based on epoch and seed
g = torch.Generator()
g.manual_seed(self.seed + self.epoch)
indices = torch.randperm(len(self.dataset), generator=g).tolist() # type: ignore[arg-type]
else: # 否则直接返回数据集长度序列
indices = list(range(len(self.dataset))) # type: ignore[arg-type]
# 是否需要补齐数据
if not self.drop_last:
# add extra samples to make it evenly divisible
padding_size = self.total_size - len(indices)
if padding_size <= len(indices):
indices += indices[:padding_size]
else:
indices += (indices * math.ceil(padding_size / len(indices)))[:padding_size]
else:
# remove tail of data to make it evenly divisible.
indices = indices[:self.total_size]
assert len(indices) == self.total_size
# subsample
# 依据自己的rank，依次返回自己的数据序号
indices = indices[self.rank:self.total_size:self.num_replicas]
assert len(indices) == self.num_samples
return iter(indices)
def __len__(self) -> int:
return self.num_samples
def set_epoch(self, epoch: int) -> None:
r"""
Sets the epoch for this sampler. When :attr:`shuffle=True`, this ensures all replicas
use a different random ordering for each epoch. Otherwise, the next iteration of this
sampler will yield the same ordering.
Args:
epoch (int): Epoch number.
"""
self.epoch = epoch
内部変数間のロジックは次のとおりです。
- データセットから長さ length を取得します。
- 構成から num_replicas (いくつかのランク) を取得し、ランク自体を取得します。
- データ セットと num_replicas の長さに応じて、num_samples と total_size を取得します。
- 最後に、インデックス = インデックス[ランク: total_size: num_replicas]; を与えます。
- インデックスを DataLoader に返す
+-----------------------------------------------------------+
| DistributedSampler |
| |
| 2 2 |
| rank +---+ num_replicas |
| + | + |
| | | | 3 |
| | | | |
| | | v |
| | | num_samples = ceil(len(dataset)/ num_replicas) |
| | | + |
| | | | |
| | | | 3 |
| | | v |
| | | total_size = num_samples * num_replicas |
| | | + |
| |4 |4 |4 |
| | | | |
| v v v |
| +-+----+------------+--------------------------------+ | +-------------+
| | | | indices | |
| | indices = indices[rank: total_size: num_replicas] +------------->+ DataLoader |
| | ^ | | 5 | |
| | | | | +-------------+
| | | | |
| +----------------------------------------------------+ |
+-----------------------------------------------------------+
|
1 | length
+------+--------+
| DataSet |
+---------------+
3.3 シャッフルデータセット
すべてのエポックはデータセットをシャッフルしますが、異なるプロセスはシャッフル後にデータセットの一貫性をどのように維持しますか?
DistributedSampler は現在のエポックを乱数シードとして使用し、インデックスを計算する前にそれを構成して、異なるプロセスが同じ乱数シードを使用するようにし、シャッフルからのデータの一貫性を確保します。
3.3.1 使用
次のコードからわかるように、分散トレーニングが必要な場合は、サンプラーのエポックを設定します。
sampler = DistributedSampler(dataset) if is_distributed else None
loader = DataLoader(dataset, shuffle=(sampler is None), ...,
sampler=sampler)
for epoch in range(start_epoch, n_epochs):
if is_distributed:
sampler.set_epoch(epoch) # 这设置epoch
train(loader)
3.3.2 パイソン
具体的にはDistributedSamplerの実装に対応します。
エポックの設定は非常に簡単です。設定するだけです。
def set_epoch(self, epoch: int) -> None:
r"""
Sets the epoch for this sampler. When :attr:`shuffle=True`, this ensures all replicas
use a different random ordering for each epoch. Otherwise, the next iteration of this
sampler will yield the same ordering.
Args:
epoch (int): Epoch number.
"""
self.epoch = epoch
ランダム シードの設定の具体的な用途は、反復関数にあります。
def __iter__(self) -> Iterator[T_co]:
if self.shuffle:
# deterministically shuffle based on epoch and seed
g = torch.Generator()
g.manual_seed(self.seed + self.epoch) # 这里设置随机种子
indices = torch.randperm(len(self.dataset), generator=g).tolist() # type: ignore[arg-type]
else:
indices = list(range(len(self.dataset))) # type: ignore[arg-type]
# 省略其他代码
3.3.3 C++
また、事前に C++ コードの DistributedRandomSampler を確認することもできます。これは C++ API であり、Python でも同じ役割を果たします。
シードとシャッフルの設定は次のようになります。
void DistributedRandomSampler::reset(optional<size_t> new_size) {
size_ = new_size.value_or(size_);
populate_indices();
std::mt19937 rand(epoch_);
std::shuffle(all_indices_.begin(), all_indices_.end(), rand);
sample_index_ = begin_index_;
}
3.3.4 まとめ
現在のロジックを次のように拡張します。
- データセットから長さ length を取得します。
- 構成、ランク自体、エポックから num_replicas (いくつかのランク) を取得します。
- エポックを使用してランダム シードを設定します。
- ランダム シードを使用してデータ セットのインデックスをスクランブルすると、スクランブルされたインデックスが常に使用されます。
- データ セットと num_replicas の長さに応じて、num_samples と total_size を取得します。
- 上記のさまざまなデータ条件と組み合わせて、最終的に indexs = index[rank: total_size: num_replicas] が与えられます。
- インデックスを DataLoader に返す
+-----------------------------------------------------------------+
| DistributedSampler |
| |
| |
| 2 3 |
| epoch +------> manual_seed(seed + epoch) +---------> indices |
| + |
| | |
| | |
| 2 2 | |
| rank +---+ num_replicas 4 | |
| + | + | |
| | | | 5 | |
| | | | | |
| | | v | |
| | | num_samples = ceil(len(dataset)/ num_replicas) | |
| | | + | |
| | | | | |
| | | | 5 | |
| | | v | |
| | | total_size = num_samples * num_replicas | |
| | | + | |
| |6 |6 |6 | |
| | | | | |
| v v v | |
| +-+----+------------+--------------------------------+ | |
| | | | |
| | indices = indices[rank: total_size: num_replicas] | <----+ |
| | ^ + | |
| | | | | |
| | | | | |
| +----------------------------------------------------+ |
+-----------------------------------------------------------------+
| |
1 | length 7 v indices
|
+-------+--------+ +-------------+
| | | |
| DataSet | | DataLoader |
| | | |
+----------------+ +-------------+
3.4 C++ のサンプラー
一部の企業は C++ で開発しているため、pytorch を緊急に使用する必要があるため、pytorch も C++ API を提供しています。次にそれを実装する方法を見てみましょう。
3.4.1 定義
そのクラスは torch\csrc\api\include\torch\data\samplers\distributed.hで定義されています。
DistributedSampler が基本クラスで、主なメンバー変数は次のとおりです。
-
size_t size_ ファイルサイズ
-
size_t num_replicas_ レプリカの数
-
size_t rank_ このサンプラーが対応するプロセスまたは GPU
-
size_t epoch このトレーニングのエポック
-
bool allow_duplicates_ バックアップを許可するかどうか
次は、DistributedRandomSampler と DistributedSequentialSampler の 2 つの派生クラスです。
/// A `Sampler` that selects a subset of indices to sample from and defines a
/// sampling behavior. In a distributed setting, this selects a subset of the
/// indices depending on the provided num_replicas and rank parameters. The
/// `Sampler` performs a rounding operation based on the `allow_duplicates`
/// parameter to decide the local sample count.
template <typename BatchRequest = std::vector<size_t>>
class DistributedSampler : public Sampler<BatchRequest> {
public:
DistributedSampler(
size_t size,
size_t num_replicas = 1,
size_t rank = 0,
bool allow_duplicates = true)
: size_(size),
num_replicas_(num_replicas),
rank_(rank),
epoch_(0),
allow_duplicates_(allow_duplicates) {}
/// Set the epoch for the current enumeration. This can be used to alter the
/// sample selection and shuffling behavior.
void set_epoch(size_t epoch) {
epoch_ = epoch;
}
size_t epoch() const {
return epoch_;
}
protected:
size_t local_sample_count() {
if (allow_duplicates_) {
return (size_ + num_replicas_ - 1) / num_replicas_;
} else {
return size_ / num_replicas_;
}
}
size_t size_;
size_t num_replicas_;
size_t rank_;
size_t epoch_;
bool allow_duplicates_;
};
/// Select samples randomly. The sampling order is shuffled at each `reset()`
/// call.
class TORCH_API DistributedRandomSampler : public DistributedSampler<> {
public:
DistributedRandomSampler(
size_t size,
size_t num_replicas = 1,
size_t rank = 0,
bool allow_duplicates = true);
/// Resets the `DistributedRandomSampler` to a new set of indices.
void reset(optional<size_t> new_size = nullopt) override;
/// Returns the next batch of indices.
optional<std::vector<size_t>> next(size_t batch_size) override;
/// Serializes the `DistributedRandomSampler` to the `archive`.
void save(serialize::OutputArchive& archive) const override;
/// Deserializes the `DistributedRandomSampler` from the `archive`.
void load(serialize::InputArchive& archive) override;
/// Returns the current index of the `DistributedRandomSampler`.
size_t index() const noexcept;
private:
void populate_indices();
size_t begin_index_;
size_t end_index_;
size_t sample_index_;
std::vector<size_t> all_indices_;
};
/// Select samples sequentially.
class TORCH_API DistributedSequentialSampler : public DistributedSampler<> {
public:
DistributedSequentialSampler(
size_t size,
size_t num_replicas = 1,
size_t rank = 0,
bool allow_duplicates = true);
/// Resets the `DistributedSequentialSampler` to a new set of indices.
void reset(optional<size_t> new_size = nullopt) override;
/// Returns the next batch of indices.
optional<std::vector<size_t>> next(size_t batch_size) override;
/// Serializes the `DistributedSequentialSampler` to the `archive`.
void save(serialize::OutputArchive& archive) const override;
/// Deserializes the `DistributedSequentialSampler` from the `archive`.
void load(serialize::InputArchive& archive) override;
/// Returns the current index of the `DistributedSequentialSampler`.
size_t index() const noexcept;
private:
void populate_indices();
size_t begin_index_;
size_t end_index_;
size_t sample_index_;
std::vector<size_t> all_indices_;
};
3.4.2 実装
クラスの特定の実装は次の場所にあります: torch\csrc\api\src\data\samplers\distributed.cpp
3.4.2.1 分散ランダムサンプラー
最初に DistributedRandomSampler を見てみましょう。
その機能は、ワーカーのrank_に従ってスクランブルされたインデックスを取得することです。各機能を論理的な順序で説明します。
- 初期化中に、シャッフルのために reset(size_) が呼び出されます。
- リセット関数の機能は、サンプラーが新しい一連のインデックスを指すようにすることです。
- 最初に populate_indices() を呼び出して、このランクに対応する開始インデックスと終了インデックスを設定します。
- populate_indices 関数では、データ インデックスが設定されます。つまり、all_indices_ が構成され、このランクの開始インデックスと終了インデックスも同時に構成されます。
- 次に all_indices_ をシャッフルします。
- 次の関数は比較的単純で、主な作業はリセットによって行われるため、この時点でデータがランダムに乱れているため、開始位置を見つけてデータの行数を返します。
以下では iota 関数を使っているので、慣れていない方もいるかもしれませんが、 iota の関数は次のとおりです。
std::vector<int> test;
test.resize(10);
std::iota(test.begin(), test.end(), 5);// 将从 5 开始的 10 次递增值赋值给 test
//test结果:5 6 7 8 9 10 11 12 13 14
具体的なコードは次のとおりです。
DistributedRandomSampler::DistributedRandomSampler(
size_t size,
size_t num_replicas,
size_t rank,
bool allow_duplicates)
: DistributedSampler(size, num_replicas, rank, allow_duplicates),
begin_index_(0),
end_index_(0),
sample_index_(0) {
// shuffle first time.
reset(size_);
}
// 每次加载新epoch时候，都要调用reset
void DistributedRandomSampler::reset(optional<size_t> new_size) {
size_ = new_size.value_or(size_);
populate_indices();
std::mt19937 rand(epoch_);
// 对于数据进行shuffle
std::shuffle(all_indices_.begin(), all_indices_.end(), rand);
sample_index_ = begin_index_;
}
void DistributedRandomSampler::populate_indices() {
size_t num_local_samples = local_sample_count();
// 得到样本数量
size_t sample_count =
num_replicas_ == 1 ? size_ : num_local_samples * num_replicas_;
all_indices_.resize(sample_count);
// std::iota 的作用是用顺序递增的值赋值指定范围内的元素
// 这里是给all_indices_设置从0开始到sample_count这些数值
std::iota(std::begin(all_indices_), std::end(all_indices_), 0);
// 如果sample count大于size_，则需要给多出来的那些index再赋一些数值
for (size_t i = size_; i < sample_count; ++i) {
// we may have added duplicate samples to make all
// replicas to have the same number of samples.
all_indices_[i] = i - size_;
}
begin_index_ = rank_ * num_local_samples; // 对应本rank的起始index
end_index_ = begin_index_ + num_local_samples; // 对应本rank的终止index
sample_index_ = begin_index_;
}
size_t DistributedRandomSampler::index() const noexcept {
return sample_index_;
}
// 注意，每次加载新epoch时候，都要调用reset，因此对于next函数来说，工作已经很小
optional<std::vector<size_t>> DistributedRandomSampler::next(
size_t batch_size) {
if (sample_index_ == end_index_) { // 已经提取完数据
return nullopt;
}
size_t end = sample_index_ + batch_size; // 本次迭代的终止位置
if (end > end_index_) {
end = end_index_;
}
auto iter = all_indices_.begin(); // 因为此时数据已经被随机打乱了，找到起始位置即可
std::vector<size_t> res(iter + sample_index_, iter + end); // 从所有数据中提取前面若干行
sample_index_ = end;
return res;
}
3.4.2.2 分散シーケンシャルサンプラー
次に、DistributedSequentialSampler を見てください。
その機能は、ワーカーのrank_に従ってオーダーのインデックスを取得することです。各機能を論理的な順序で説明します。
- リセット関数ははるかに簡単です。 populate_indices を使用してインデックスを順番に設定するだけです。
- 次の関数は比較的複雑で、インデックスを順番に返すだけでなく、次の開始位置を設定します。
DistributedSequentialSampler::DistributedSequentialSampler(
size_t size,
size_t num_replicas,
size_t rank,
bool allow_duplicates)
: DistributedSampler(size, num_replicas, rank, allow_duplicates),
begin_index_(0),
end_index_(0),
sample_index_(0) {
populate_indices(); // 这里会设定本rank对应的起始位置
}
void DistributedSequentialSampler::reset(optional<size_t> new_size) {
size_t size = new_size.value_or(size_);
if (size != size_) {
size_ = size;
populate_indices();
} else {
sample_index_ = begin_index_;
}
}
void DistributedSequentialSampler::populate_indices() {
begin_index_ = rank_ * local_sample_count(); // 本rank对应的起始位置
end_index_ = begin_index_ + local_sample_count();
sample_index_ = begin_index_;
}
size_t DistributedSequentialSampler::index() const noexcept {
return sample_index_;
}
optional<std::vector<size_t>> DistributedSequentialSampler::next(
size_t batch_size) {
if (sample_index_ == end_index_) { // 已经循环结束
return nullopt;
}
size_t end = sample_index_ + batch_size; // 本次的终止行
if (end > end_index_) {
end = end_index_;
}
std::vector<size_t> res(end - sample_index_); // 返回的vector大小
// 给res设置从sample_index_开始递增(end - sample_index_)这么大的这些数值，这就是顺序返回了index
std::iota(std::begin(res), std::end(res), sample_index_);
if (end >= size_) {
for (size_t& index : res) { //遍历 vector，得到本次的index
index = index % size_;
}
}
sample_index_ = end; // 设置下次开始行
return res;
}
0xFF リファレンス
畳み込みニューラル ネットワークの並列化モデル -- 畳み込みニューラル ネットワークを並列化するための奇妙なトリック
torch.utils.data の PyTorch ソース コード解釈: データ処理の全プロセスを分析する
大規模機械学習の分野に対する理解と認識について教えてください。
pytorch (分散) データ並列の個人的な実践のまとめ - DataParallel/DistributedDataParallel