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

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))
Logo

有“AI”的1024 = 2048,欢迎大家加入2048 AI社区

更多推荐