参考:https://www.jianshu.com/p/aee6a3d72014
# 网络,netg为生成器,netd为判别器
netg, netd = NetG(opt), NetD(opt)
# 把所有的张量加载到CPU中
map_location = lambda storage, loc: storage
# 把所有的张量加载到GPU 1中
#torch.load('tensors.pt', map_location=lambda storage, loc: storage.cuda(1))
#也可以写成:
#device = torch.device('cpu')
#netd.load_state_dict(t.load(opt.netd_path, map_location=device))
#或:
#netd.load_state_dict(t.load(opt.netd_path))
#netd.to(device)
if opt.netd_path: #是否指定训练好的预训练模型,加载模型参数
netd.load_state_dict(t.load(opt.netd_path, map_location=map_location))
if opt.netg_path:
netg.load_state_dict(t.load(opt.netg_path, map_location=map_location))
netd.to(device)
netg.to(device)
# 定义优化器和损失,学习率都默认为2e-4,beta1默认为0.5
optimizer_g = t.optim.Adam(netg.parameters(), opt.lr1, betas=(opt.beta1, 0.999))
optimizer_d = t.optim.Adam(netd.parameters(), opt.lr2, betas=(opt.beta1, 0.999))
criterion = t.nn.BCELoss().to(device)
# 真图片label为1,假图片label为0
# noises为生成网络的输入
true_labels = t.ones(opt.batch_size).to(device)
fake_labels = t.zeros(opt.batch_size).to(device)
fix_noises = t.randn(opt.batch_size, opt.nz, 1, 1).to(device)#opt.nz为噪声维度,默认为100
noises = t.randn(opt.batch_size, opt.nz, 1, 1).to(device)
#AverageValueMeter测量并返回添加到其中的任何数字集合的平均值和标准差,
#对度量一组示例的平均损失是有用的。
errord_meter = AverageValueMeter()
errorg_meter = AverageValueMeter()
再加载预训练模型时,最好指定map_location。因为如果程序之前在GPU上运行,那么模型就会被存为torch.cuda.Tensor,这样加载时会默认将数据加载至显存。如果运行该程序的计算机中没有GPU,加载就会报错,故通过指定map_location将Tensor默认加载入内存(CPU),待有需要再移至显存