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

Update train.py 

solve the Multi-GPU --resume Error #138
https://github.com/ultralytics/yolov3/issues/138
This commit is contained in:
perry0418 2019-03-21 12:02:57 +08:00 committed by GitHub
parent 2cd6805063
commit 35396adc9c
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
train.py
View File

@ -32,7 +32,7 @@ def train(
train_path = parse_data_cfg(data_cfg)['train']
# Initialize model
model = Darknet(cfg, img_size)
model = Darknet(cfg, img_size).to(device)
# Get dataloader
dataloader = LoadImagesAndLabels(train_path, batch_size, img_size, augment=True)
@ -43,7 +43,7 @@ def train(
start_epoch = 0
best_loss = float('inf')
if resume:
checkpoint = torch.load(latest, map_location='cpu')
checkpoint = torch.load(latest, map_location=device)
# Load weights to resume from
model.load_state_dict(checkpoint['model'])