zjc
/
yolov5
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

Merge pull request #143 from Lornatang/Fix-multi-GPU-training-errors 

Fix DDP bug in single process multiple device use cases
This commit is contained in:
Glenn Jocher 2020-06-20 11:10:40 -07:00 committed by GitHub
parent 024a42f4fd 899f1d4bde
commit ba340ed72d
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
requirements.txt
View File

@ -2,7 +2,7 @@
Cython Cython
numpy==1.17 numpy==1.17
opencv-python opencv-python
torch>=1.5 torch==1.4
matplotlib matplotlib
pillow pillow
tensorboard tensorboard

1
train.py
View File

@ -151,6 +151,7 @@ def train(hyp):
world_size=1, # number of nodes world_size=1, # number of nodes
rank=0) # node rank rank=0) # node rank
model = torch.nn.parallel.DistributedDataParallel(model) model = torch.nn.parallel.DistributedDataParallel(model)
# pip install torch==1.4.0+cu100 torchvision==0.5.0+cu100 -f https://download.pytorch.org/whl/torch_stable.html
# Dataset # Dataset
dataset = LoadImagesAndLabels(train_path, imgsz, batch_size, dataset = LoadImagesAndLabels(train_path, imgsz, batch_size,