pytorch-YOLOv4
pytorch-YOLOv4 copied to clipboard
Does it support multi GPUs?
How to train it with numlti gpus?
Sure it does. Check README.md
Train
you can set parameters in cfg.py.
python train.py -g [GPU_ID] -dir [Dataset direction] ...
python train.py -l 0.001 -g 0,1 -pretrained ./yolov4.pth -classes 12 -dir ../VisDrone2019-DET -train_label_path ./data/train.txt the 0,1 is the id of GPU,it means you can train this model by the first and second GPU of your computer
I use multi GPUs find that the result is error when i use models.py test images,the error is :
Traceback (most recent call last):
File "models.py", line 487, in
when i use strict=False in model.loas_state_dict(),the result is null .could you help me?
I use multi GPUs find that the result is error when i use models.py test images,the error is : Traceback (most recent call last): File "models.py", line 487, in model.load_state_dict(pretrained_dict) File "/home/ubuntu/anaconda3/envs/yolov4/lib/python3.6/site-packages/torch/nn/modules/module.py", line 1052, in load_state_dict self.class.name, "\n\t".join(error_msgs))) RuntimeError: Error(s) in loading state_dict for Yolov4: Missing key(s) in state_dict:"down1.conv1.conv.0.weight", "down1.conv1.conv.1.weight", "down1.conv1.conv.1.bias", "down1.conv1.conv.1.running_mean"...
when i use strict=False in model.loas_state_dict(),the result is null .could you help me?
Yes, same problem when I use strict=False
Hey, I found a way to properly load the model (don't ask me why it works):
model = Yolov4(yolov4conv137weight=None, n_classes=1, inference=True)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
if torch.cuda.device_count() > 1:
model = torch.nn.DataParallel(model)
model.to(device=device)
pretrained_dict = torch.load(model_path, map_location=torch.device('cuda'))
model.load_state_dict(pretrained_dict)
model.cuda()