perf(model): 在cpu中保存模型权重

Author: zjZSTU

py/lib/train.py

@@ -89,11 +89,11 @@ def train_model(data_loader, model, criterion, optimizer, lr_scheduler, num_epoc
             # deep copy the model
             if epoch_loss < best_loss:
                 best_loss = epoch_loss
-                best_model_weights = copy.deepcopy(model.state_dict())
+                best_model_weights = copy.deepcopy(model.cpu().state_dict())
+                model = model.to(device)
 
                 file.check_dir('../models')
-                # util.save_checkpoint('../models/checkpoint_yolo_v1_%d.pth' % (epoch), epoch, model, optimizer, loss)
-                file.save_model(model, '../models/checkpoint_yolo_v1_%d.pth' % (epoch))
+                file.save_model(best_model_weights, '../models/checkpoint_yolo_v1_%d.pth' % (epoch))
                 print('save model')
 
         print()

py/lib/utils/file.py

@@ -72,8 +72,8 @@ def parse_location_xml(xml_path):
         return np.array(bndboxs), name_list
 
 
-def save_model(model, model_save_path):
-    torch.save(model.state_dict().cpu(), model_save_path)
+def save_model(model_weights, model_save_path):
+    torch.save(model_weights, model_save_path)
 
 
 def save_checkpoint(model_save_path, epoch, model, optimizer, loss):