updates
This commit is contained in:
parent
0fb4a46ace
commit
d5815ebfd2
8
train.py
8
train.py
|
@ -132,6 +132,10 @@ def train():
|
||||||
# possible weights are '*.weights', 'yolov3-tiny.conv.15', 'darknet53.conv.74' etc.
|
# possible weights are '*.weights', 'yolov3-tiny.conv.15', 'darknet53.conv.74' etc.
|
||||||
load_darknet_weights(model, weights)
|
load_darknet_weights(model, weights)
|
||||||
|
|
||||||
|
# Mixed precision training https://github.com/NVIDIA/apex
|
||||||
|
if mixed_precision:
|
||||||
|
model, optimizer = amp.initialize(model, optimizer, opt_level='O1', verbosity=0)
|
||||||
|
|
||||||
# Scheduler https://github.com/ultralytics/yolov3/issues/238
|
# Scheduler https://github.com/ultralytics/yolov3/issues/238
|
||||||
# lf = lambda x: 1 - x / epochs # linear ramp to zero
|
# lf = lambda x: 1 - x / epochs # linear ramp to zero
|
||||||
# lf = lambda x: 10 ** (hyp['lrf'] * x / epochs) # exp ramp
|
# lf = lambda x: 10 ** (hyp['lrf'] * x / epochs) # exp ramp
|
||||||
|
@ -153,10 +157,6 @@ def train():
|
||||||
# plt.tight_layout()
|
# plt.tight_layout()
|
||||||
# plt.savefig('LR.png', dpi=300)
|
# plt.savefig('LR.png', dpi=300)
|
||||||
|
|
||||||
# Mixed precision training https://github.com/NVIDIA/apex
|
|
||||||
if mixed_precision:
|
|
||||||
model, optimizer = amp.initialize(model, optimizer, opt_level='O1', verbosity=0)
|
|
||||||
|
|
||||||
# Initialize distributed training
|
# Initialize distributed training
|
||||||
if device.type != 'cpu' and torch.cuda.device_count() > 1:
|
if device.type != 'cpu' and torch.cuda.device_count() > 1:
|
||||||
dist.init_process_group(backend='nccl', # 'distributed backend'
|
dist.init_process_group(backend='nccl', # 'distributed backend'
|
||||||
|
|
Loading…
Reference in New Issue