ソースを参照

Remove DDP MultiHeadAttention fix (#3768)

modifyDataloader
Glenn Jocher GitHub 3年前
コミット
f2d97ebb25
この署名に対応する既知のキーがデータベースに存在しません GPGキーID: 4AEE18F83AFDEB23
1個のファイルの変更1行の追加3行の削除
  1. +1
    -3
      train.py

+ 1
- 3
train.py ファイルの表示

@@ -252,9 +252,7 @@ def train(hyp, # path/to/hyp.yaml or hyp dictionary

# DDP mode
if cuda and RANK != -1:
model = DDP(model, device_ids=[LOCAL_RANK], output_device=LOCAL_RANK,
# nn.MultiheadAttention incompatibility with DDP https://github.com/pytorch/pytorch/issues/26698
find_unused_parameters=any(isinstance(layer, nn.MultiheadAttention) for layer in model.modules()))
model = DDP(model, device_ids=[LOCAL_RANK], output_device=LOCAL_RANK)

# Model parameters
hyp['box'] *= 3. / nl # scale to layers

読み込み中…
キャンセル
保存