参考：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)，待有需要再移至显存