def train(opt):
print("Training model with the following parameters:")
print("\t number of stages: {}".format(opt.train_stages))
print("\t number of concurrently trained stages: {}".format(opt.train_depth))
print("\t learning rate scaling: {}".format(opt.lr_scale))
print("\t non-linearity: {}".format(opt.activation))
real, real2 = functions.read_two_domains(opt)
# real = functions.read_image(opt)
# print(0, real.shape)
real = functions.adjust_scales2image(real, opt)
reals = functions.create_reals_pyramid(real, opt)
real2 = functions.adjust_scales2image(real2, opt)
reals2 = functions.create_reals_pyramid(real2, opt)
generator, generator2 = init_G(opt)
fixed_noise = []
noise_amp = []
fixed_noise2 = []
noise_amp2 = []
for scale_num in range(opt.stop_scale+1):
opt.out_ = functions.generate_dir2save(opt)
opt.outf = '%s/%d' % (opt.out_,scale_num)
try:
os.makedirs(opt.outf)
except OSError:
print(OSError)
pass
functions.save_image('{}/real_scale.jpg'.format(opt.outf), reals[scale_num])
d_curr, d_curr2 = init_D(opt)
if scale_num > 0:
d_curr.load_state_dict(torch.load('%s/%d/netD.pth' % (opt.out_,scale_num-1)))
generator.init_next_stage()
d_curr2.load_state_dict(torch.load('%s/%d/netD2.pth' % (opt.out_,scale_num-1)))
generator2.init_next_stage()
writer = SummaryWriter(log_dir=opt.outf)
fixed_noise, noise_amp, generator, d_curr, fixed_noise2, noise_amp2, generator2, d_curr2 = \
train_single_scale(d_curr, generator, reals, fixed_noise, noise_amp, d_curr2, generator2, reals2,
fixed_noise2, noise_amp2, opt, scale_num, writer)
torch.save(fixed_noise, '%s/fixed_noise.pth' % (opt.out_))
torch.save(generator, '%s/G.pth' % (opt.out_))
torch.save(reals, '%s/reals.pth' % (opt.out_))
torch.save(noise_amp, '%s/noise_amp.pth' % (opt.out_))
torch.save(fixed_noise2, '%s/fixed_noise2.pth' % (opt.out_))
torch.save(generator2, '%s/G2.pth' % (opt.out_))
torch.save(reals2, '%s/reals2.pth' % (opt.out_))
torch.save(noise_amp2, '%s/noise_amp2.pth' % (opt.out_))
del d_curr, d_curr2
writer.close()
return
def train(opt):
print("Training model with the following parameters:")
print("\t number of stages: {}".format(opt.train_stages))
print("\t number of concurrently trained stages: {}".format(
opt.train_depth))
print("\t learning rate scaling: {}".format(opt.lr_scale))
print("\t non-linearity: {}".format(opt.activation))
real = functions.read_image(opt)
real = functions.adjust_scales2image(real, opt)
reals = functions.create_reals_pyramid(real, opt)
print("Training on image pyramid: {}".format([r.shape for r in reals]))
print("")
if opt.naive_img != "":
naive_img = functions.read_image_dir(opt.naive_img, opt)
naive_img_large = imresize_to_shape(naive_img, reals[-1].shape[2:],
opt)
naive_img = imresize_to_shape(naive_img, reals[0].shape[2:], opt)
naive_img = functions.convert_image_np(naive_img) * 255.0
else:
naive_img = None
naive_img_large = None
if opt.fine_tune:
img_to_augment = naive_img
else:
img_to_augment = functions.convert_image_np(reals[0]) * 255.0
if opt.train_mode == "editing":
opt.noise_scaling = 0.1
generator = init_G(opt)
if opt.fine_tune:
for _ in range(opt.train_stages - 1):
generator.init_next_stage()
generator.load_state_dict(
torch.load(
'{}/{}/netG.pth'.format(opt.model_dir, opt.train_stages - 1),
map_location="cuda:{}".format(torch.cuda.current_device())))
fixed_noise = []
noise_amp = []
for scale_num in range(opt.start_scale, opt.train_stages):
opt.out_ = functions.generate_dir2save(opt)
opt.outf = '%s/%d' % (opt.out_, scale_num)
try:
os.makedirs(opt.outf)
except OSError:
print(OSError)
pass
functions.save_image('{}/real_scale.jpg'.format(opt.outf),
reals[scale_num])
d_curr = init_D(opt)
if opt.fine_tune:
d_curr.load_state_dict(
torch.load('{}/{}/netD.pth'.format(opt.model_dir,
opt.train_stages - 1),
map_location="cuda:{}".format(
torch.cuda.current_device())))
elif scale_num > 0:
d_curr.load_state_dict(
torch.load('%s/%d/netD.pth' % (opt.out_, scale_num - 1)))
generator.init_next_stage()
writer = SummaryWriter(log_dir=opt.outf)
fixed_noise, noise_amp, generator, d_curr = train_single_scale(
d_curr, generator, reals, img_to_augment, naive_img,
naive_img_large, fixed_noise, noise_amp, opt, scale_num, writer)
torch.save(fixed_noise, '%s/fixed_noise.pth' % (opt.out_))
torch.save(generator, '%s/G.pth' % (opt.out_))
torch.save(reals, '%s/reals.pth' % (opt.out_))
torch.save(noise_amp, '%s/noise_amp.pth' % (opt.out_))
del d_curr
writer.close()
return