Browse Source

backbone as FP16, save default to FP32

5.0
Glenn Jocher 4 years ago
parent
commit
cce95e744d
2 changed files with 2 additions and 3 deletions
  1. +1
    -1
      train.py
  2. +1
    -2
      utils/utils.py

+ 1
- 1
train.py View File

ckpt = {'epoch': epoch, ckpt = {'epoch': epoch,
'best_fitness': best_fitness, 'best_fitness': best_fitness,
'training_results': f.read(), 'training_results': f.read(),
'model': ema.ema.module.half() if hasattr(model, 'module') else ema.ema.half(),
'model': ema.ema.module if hasattr(model, 'module') else ema.ema,
'optimizer': None if final_epoch else optimizer.state_dict()} 'optimizer': None if final_epoch else optimizer.state_dict()}


# Save last, best and delete # Save last, best and delete

+ 1
- 2
utils/utils.py View File

def create_backbone(f='weights/best.pt', s='weights/backbone.pt'): # from utils.utils import *; create_backbone() def create_backbone(f='weights/best.pt', s='weights/backbone.pt'): # from utils.utils import *; create_backbone()
# create backbone 's' from 'f' # create backbone 's' from 'f'
device = torch.device('cpu') device = torch.device('cpu')
x = torch.load(f, map_location=device)
torch.save(x, s) # update model if SourceChangeWarning
x = torch.load(s, map_location=device) x = torch.load(s, map_location=device)


x['optimizer'] = None x['optimizer'] = None
x['training_results'] = None x['training_results'] = None
x['epoch'] = -1 x['epoch'] = -1
x['model'].half() # to FP16
for p in x['model'].parameters(): for p in x['model'].parameters():
p.requires_grad = True p.requires_grad = True
torch.save(x, s) torch.save(x, s)

Loading…
Cancel
Save