From b88c4568bae2de03bda72ceb2e63a7d1f2662358 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Sat, 24 Aug 2019 12:37:55 +0200 Subject: [PATCH] updates --- utils/utils.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/utils/utils.py b/utils/utils.py index a3c9520e..5932c8b8 100755 --- a/utils/utils.py +++ b/utils/utils.py @@ -562,11 +562,15 @@ def print_model_biases(model): 'classification: %.2f+/-%.2f' % (b[:, 5:].mean(), b[:, 5:].std())) -def strip_optimizer_from_checkpoint(filename='weights/best.pt'): +def strip_optimizer(f='weights/best.pt'): # Strip optimizer from *.pt files for lighter files (reduced by 2/3 size) - a = torch.load(filename, map_location='cpu') - a['optimizer'] = [] - torch.save(a, filename.replace('.pt', '_lite.pt')) + x = torch.load(f) + x['optimizer'] = None + # x['training_results'] = None + # x['epoch'] = -1 + # for p in x['model']: + # p.requires_grad = True + torch.save(x, f) def coco_class_count(path='../coco/labels/train2014/'):