Bienqq
/
car-detection-bayes
2
1
Fork 0
Code Issues Pull Requests Releases Wiki Activity

updates

This commit is contained in:
Glenn Jocher 2018-11-29 11:59:29 +01:00
parent d5331be0a0
commit af0033c9e9
1 changed files with 6 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
train.py
View File

@ -58,8 +58,9 @@ def main(opt):
model.load_state_dict(checkpoint['model']) model.load_state_dict(checkpoint['model'])
if torch.cuda.device_count() > 1: if torch.cuda.device_count() > 1:
print('Using ', torch.cuda.device_count(), ' GPUs') raise Exception('Multi-GPU not currently supported: https://github.com/ultralytics/yolov3/issues/21')
model = nn.DataParallel(model) # print('Using ', torch.cuda.device_count(), ' GPUs')
# model = nn.DataParallel(model)
model.to(device).train() model.to(device).train()
# # Transfer learning (train only YOLO layers) # # Transfer learning (train only YOLO layers)
@ -87,8 +88,9 @@ def main(opt):
load_weights(model, 'weights/darknet53.conv.74') load_weights(model, 'weights/darknet53.conv.74')
if torch.cuda.device_count() > 1: if torch.cuda.device_count() > 1:
print('Using ', torch.cuda.device_count(), ' GPUs') raise Exception('Multi-GPU not currently supported: https://github.com/ultralytics/yolov3/issues/21')
model = nn.DataParallel(model) # print('Using ', torch.cuda.device_count(), ' GPUs')
# model = nn.DataParallel(model)
model.to(device).train() model.to(device).train()
# Set optimizer # Set optimizer