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

Update train.py

This commit is contained in:
Glenn Jocher 2019-03-25 15:05:35 +01:00 committed by GitHub
parent c7192f64c9
commit 06d264198c
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 0 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
train.py
View File

@ -61,11 +61,6 @@ def train(
elif cfg.endswith('yolov3-tiny.cfg'):
cutoff = load_darknet_weights(model, weights + 'yolov3-tiny.conv.15')
#initialize for distributed training
if torch.cuda.device_count() > 1:
dist.init_process_group(backend=opt.backend, init_method=opt.dist_url,world_size=opt.world_size, rank=opt.rank)
model = torch.nn.parallel.DistributedDataParallel(model)
# Transfer learning (train only YOLO layers)
# for i, (name, p) in enumerate(model.named_parameters()):
# p.requires_grad = True if (p.shape[0] == 255) else False