pytorch：模型的保存与导出

方法一： 保存模型和模型参数

torch.save( network, savePath )

  def save_network( save_dir, network, network_label, epoch_label):
      save_filename = '%s_net_%s.pth' % (epoch_label, network_label)
      save_path = os.path.join(save_dir, save_filename)
      torch.save(network, save_path) 

network = torch.load( loadPath )

  def load( load_dir, network_label, epoch_label='latest'):
      load_filename = '%s_net_%s.pth' % (epoch_label, network_label)
      load_path = os.path.join(load_dir, load_filename)
      return torch.load(load_path )

特点：模型的导入和导出很容易，但是如果数据量比较大，会消耗大量时间。占用的内存也比较高。

方法二： 只保存模型参数 （推荐）

torch.save( network.state_dict(), savePath )

  def save_network(save_dir, network, network_label, epoch_label):
      save_filename = '%s_net_%s_params.pkl' % (epoch_label, network_label)
      save_path = os.path.join(save_dir, save_filename)
      torch.save(network.state_dict(), save_path)  

由于模型保存的是参数，所以在测试阶段，要先定义网络，再把导出模型参数赋值给定义的网络：

# Load model
   #定义网络
   G = Generator(opt.input_nc, opt.output_nc)
   G.cuda()
   # 把保存的参数导出，并赋值给网络
   model_dir = os.path.join(opt.checkpoints_dir, opt.name)
   load_filename = '%s_net_%s_params.pkl' % (epoch_label, network_label)
   load_path = os.path.join(save_dir, load_filename)
   G.load_state_dict(torch.load(load_path))