def train(epochs, iterations, batchsize, validsize, outdir, modeldir, extension, train_size, valid_size, data_path, sketch_path, digi_path, learning_rate, beta1, weight_decay): # Dataset definition dataset = DatasetLoader(data_path, sketch_path, digi_path, extension, train_size, valid_size) print(dataset) x_val, t_val = dataset.valid(validsize) # Model & Optimizer definition unet = UNet() unet.to_gpu() unet_opt = set_optimizer(unet, learning_rate, beta1, weight_decay) discriminator = Discriminator() discriminator.to_gpu() dis_opt = set_optimizer(discriminator, learning_rate, beta1, weight_decay) # Loss function definition lossfunc = Pix2pixLossCalculator() # Visualization definition visualizer = Visualizer() for epoch in range(epochs): sum_dis_loss = 0 sum_gen_loss = 0 for batch in range(0, iterations, batchsize): x, t = dataset.train(batchsize) # Discriminator update y = unet(x) y.unchain_backward() dis_loss = lossfunc.dis_loss(discriminator, y, t) discriminator.cleargrads() dis_loss.backward() dis_opt.update() sum_dis_loss += dis_loss.data # Generator update y = unet(x) gen_loss = lossfunc.gen_loss(discriminator, y) gen_loss += lossfunc.content_loss(y, t) unet.cleargrads() gen_loss.backward() unet_opt.update() sum_gen_loss += gen_loss.data if batch == 0: serializers.save_npz(f"{modeldir}/unet_{epoch}.model", unet) with chainer.using_config("train", False): y = unet(x_val) x = x_val.data.get() t = t_val.data.get() y = y.data.get() visualizer(x, t, y, outdir, epoch, validsize) print(f"epoch: {epoch}") print( f"dis loss: {sum_dis_loss/iterations} gen loss: {sum_gen_loss/iterations}" )
test_path="./test.png" test=prepare_image(test_path) test=chainer.as_variable(xp.array(test).astype(xp.float32)) test=F.tile(test,(framesize,1,1,1)) image_encoder=ImageEncoder() image_encoder.to_gpu() enc_opt=set_optimizer(image_encoder) key_point_detector = KeyPointDetector() key_point_detector.to_gpu() key_opt = set_optimizer(key_point_detector) making_optical_flow = UNet(in_ch=4) making_optical_flow.to_gpu() ref_opt=set_optimizer(making_optical_flow) generator = Generator(in_ch=3) generator.to_gpu() gen_opt = set_optimizer(generator) discriminator_temporal = Discriminator_temporal() discriminator_temporal.to_gpu() dis_temp_opt=set_optimizer(discriminator_temporal) discriminator_image = Discriminator_image() discriminator_image.to_gpu() dis_img_opt = set_optimizer(discriminator_image) for epoch in range(epochs):
gen_opt_xy = set_optimizer(generator_xy) generator_yx = Generator() generator_yx.to_gpu() gen_opt_yx = set_optimizer(generator_yx) discriminator_xy = Discriminator() discriminator_xy.to_gpu() dis_opt_xy = set_optimizer(discriminator_xy) discriminator_yx = Discriminator() discriminator_yx.to_gpu() dis_opt_yx = set_optimizer(discriminator_yx) predictor_x = UNet() predictor_x.to_gpu() pre_opt_x = set_optimizer(predictor_x) predictor_y = UNet() predictor_y.to_gpu() pre_opt_y = set_optimizer(predictor_y) for epoch in range(epochs): sum_dis_loss = 0 sum_gen_loss = 0 for batch in range(0, iterations, batchsize): x_box = [] y_box = [] rnd1 = np.random.randint(x_len - batchsize) rnd2 = np.random.randint(y_len - batchsize) for index in range(batchsize):