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