jasder
/
yolov5
1
0
Fork 0
Code Issues Pull Requests 1 Releases Wiki Activity

FP16 EMA bug fix

Browse Source
pull/1/head
Glenn Jocher 5 years ago
parent 2b1892430f
commit 2377e5f6ee
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File

4
utils/torch_utils.py
Unescape View File

@ -195,8 +195,8 @@ class ModelEMA:
def __init__(self, model, decay=0.9999, updates=0): def __init__(self, model, decay=0.9999, updates=0):
# Create EMA # Create EMA
self.ema = deepcopy(model.module if is_parallel(model) else model).eval() # FP32 EMA self.ema = deepcopy(model.module if is_parallel(model) else model).eval() # FP32 EMA
if next(model.parameters()).device.type != 'cpu': # if next(model.parameters()).device.type != 'cpu':
self.ema.half() # FP16 EMA # self.ema.half() # FP16 EMA
self.updates = updates # number of EMA updates self.updates = updates # number of EMA updates
self.decay = lambda x: decay * (1 - math.exp(-x / 2000)) # decay exponential ramp (to help early epochs) self.decay = lambda x: decay * (1 - math.exp(-x / 2000)) # decay exponential ramp (to help early epochs)
for p in self.ema.parameters(): for p in self.ema.parameters():

Write Preview
Loading…
Cancel
Save