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