def load_checkpoint(opts):
"""Loads the generator and discriminator models from checkpoints.
"""
G_XtoY_path = os.path.join(opts.load, 'G_XtoY.pkl')
G_YtoX_path = os.path.join(opts.load, 'G_YtoX.pkl')
D_X_path = os.path.join(opts.load, 'D_X.pkl')
D_Y_path = os.path.join(opts.load, 'D_Y.pkl')
G_XtoY = CycleGenerator(conv_dim=opts.g_conv_dim,
init_zero_weights=opts.init_zero_weights)
G_YtoX = CycleGenerator(conv_dim=opts.g_conv_dim,
init_zero_weights=opts.init_zero_weights)
D_X = DCDiscriminator(conv_dim=opts.d_conv_dim)
D_Y = DCDiscriminator(conv_dim=opts.d_conv_dim)
G_XtoY.load_state_dict(
torch.load(G_XtoY_path, map_location=lambda storage, loc: storage))
G_YtoX.load_state_dict(
torch.load(G_YtoX_path, map_location=lambda storage, loc: storage))
D_X.load_state_dict(
torch.load(D_X_path, map_location=lambda storage, loc: storage))
D_Y.load_state_dict(
torch.load(D_Y_path, map_location=lambda storage, loc: storage))
if torch.cuda.is_available():
G_XtoY.cuda()
G_YtoX.cuda()
D_X.cuda()
D_Y.cuda()
print('Models moved to GPU.')
return G_XtoY, G_YtoX, D_X, D_Y
def create_model(opts):
G_XtoY = CycleGenerator(init_zero_weights=opts.init_zero_weights)
G_YtoX = CycleGenerator(init_zero_weights=opts.init_zero_weights)
D_X = PatchGANDiscriminator()
D_Y = PatchGANDiscriminator()
if torch.cuda.is_available():
G_XtoY.cuda()
G_YtoX.cuda()
D_X.cuda()
D_Y.cuda()
print('Models moved to GPU.')
return G_XtoY, G_YtoX, D_X, D_Y
def create_model(opts):
"""Builds the generators and discriminators.
"""
G_XtoY = CycleGenerator(conv_dim=opts.g_conv_dim,
init_zero_weights=opts.init_zero_weights)
G_YtoX = CycleGenerator(conv_dim=opts.g_conv_dim,
init_zero_weights=opts.init_zero_weights)
D_X = DCDiscriminator(conv_dim=opts.d_conv_dim)
D_Y = DCDiscriminator(conv_dim=opts.d_conv_dim)
if torch.cuda.is_available():
G_XtoY.cuda()
G_YtoX.cuda()
D_X.cuda()
D_Y.cuda()
return G_XtoY, G_YtoX, D_X, D_Y
def create_model(opts):
"""Builds the generators and discriminators.
"""
G_XtoY = CycleGenerator(conv_dim=opts.g_conv_dim, init_zero_weights=opts.init_zero_weights, batch_norm=not opts.disable_bn)
G_YtoX = CycleGenerator(conv_dim=opts.g_conv_dim, init_zero_weights=opts.init_zero_weights, batch_norm=not opts.disable_bn)
D_X = DCDiscriminator(conv_dim=opts.d_conv_dim)
D_Y = DCDiscriminator(conv_dim=opts.d_conv_dim)
print_models(G_XtoY, G_YtoX, D_X, D_Y)
if torch.cuda.is_available():
G_XtoY.cuda()
G_YtoX.cuda()
D_X.cuda()
D_Y.cuda()
print('Models moved to GPU.')
return G_XtoY, G_YtoX, D_X, D_Y
opt.n_epochs = 100
opt.n_epochs_decay = 100
opt.port = 35850
opt.test = 1
print(opt)
if torch.cuda.is_available() and not opt.cuda:
print("WARNING: You have a CUDA device, so you should probably run with --cuda")
###### Definition of variables ######
# Networks
netG_A2B = CycleGenerator(opt.A_nc, opt.B_nc)
netG_B2A = CycleGenerator(opt.B_nc, opt.A_nc)
if opt.cuda:
netG_A2B.cuda()
netG_B2A.cuda()
# Load state dicts
netG_A2B.load_state_dict(torch.load(opt.generator_A2B))
netG_B2A.load_state_dict(torch.load(opt.generator_B2A))
# Set model's test mode
netG_A2B.eval()
netG_B2A.eval()
# Inputs & targets memory allocation
Tensor = torch.cuda.FloatTensor if opt.cuda else torch.Tensor
input_A = Tensor(opt.batch_size, opt.A_nc, opt.size, opt.size)
input_B = Tensor(opt.batch_size, opt.B_nc, opt.size, opt.size)