方法一: 保存模型和模型参数
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))