def load_network():
# os.environ["CUDA_VISIBLE_DEVICES"] = cfg.GPU_ID
# os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"
print('###################################')
print("##### Build Network #####")
print('###################################')
nets = []
netG = network.Res_Generator(cfg.TRAIN.ngf, cfg.TRAIN.num_resblock)
netG = nn.DataParallel(netG)
# netG.to(device)
netD = network.Patch_Discriminator(cfg.TRAIN.ndf)
netD = nn.DataParallel(netD)
# netD.to(device)
nets.append(netG)
nets.append(netD)
# print_networks(nets, debug=cfg.DEBUG)
for net in nets:
net.cuda()
return nets
def load_network(model_path):
os.environ["CUDA_VISIBLE_DEVICES"] = cfg.TEST.GPU_ID
print('###################################')
print("##### Load Network #####")
print('###################################')
nets = []
netG = network.Res_Generator(cfg.TRAIN.ngf, cfg.TRAIN.num_resblock)
netG.load_state_dict(torch.load(model_path)['state_dict'])
nets.append(netG)
for net in nets:
net.cuda()
print('Finished !')
return nets
def load_network():
os.environ["CUDA_VISIBLE_DEVICES"] = cfg.GPU_ID
print ('###################################')
print ("##### Build Network #####")
print ('###################################')
netG = network.Res_Generator(ngf=64, nz=(2048+50))
netD = network.DC_Discriminator(ndf=64)
netRN = network.ResNet50()
netE = network.Ensemble(netRN, netG)
nets = []
nets.append(netE)
nets.append(netD)
# print_networks(nets, debug=True)
for net in nets:
net.cuda()
return nets
