From a5601a1afb4586b3a4fb9ba170c7cdf3b0e68d14 Mon Sep 17 00:00:00 2001 From: root Date: Mon, 2 Nov 2020 09:22:19 +0530 Subject: [PATCH 1/3] Converted to TF2 --- test_code/cartoonize.py | 13 +- test_code/guided_filter.py | 20 +- test_code/network.py | 14 +- test_code/saved_models/checkpoint | 0 .../model-33999.data-00000-of-00001 | Bin test_code/saved_models/model-33999.index | Bin test_code/test_images/actress2.jpg | Bin test_code/test_images/china6.jpg | Bin test_code/test_images/food16.jpg | Bin test_code/test_images/food6.jpg | Bin test_code/test_images/liuyifei4.jpg | Bin test_code/test_images/london1.jpg | Bin test_code/test_images/mountain4.jpg | Bin test_code/test_images/mountain5.jpg | Bin test_code/test_images/national_park1.jpg | Bin test_code/test_images/party5.jpg | Bin test_code/test_images/party7.jpg | Bin train_code/guided_filter.py | 4 +- train_code/layers.py | 185 +++++---- train_code/loss.py | 381 +++++++++--------- train_code/network.py | 29 +- train_code/pretrain.py | 34 +- .../__pycache__/__init__.cpython-37.pyc | Bin 0 -> 267 bytes .../__pycache__/core.cpython-37.pyc | Bin 0 -> 3564 bytes .../__pycache__/measure.cpython-37.pyc | Bin 0 -> 3854 bytes .../__pycache__/structure.cpython-37.pyc | Bin 0 -> 4319 bytes .../__pycache__/util.cpython-37.pyc | Bin 0 -> 2012 bytes train_code/selective_search/core.py | 4 +- train_code/selective_search/structure.py | 12 +- train_code/train.py | 88 ++-- train_code/utils.py | 1 - 31 files changed, 392 insertions(+), 393 deletions(-) mode change 100644 => 100755 test_code/saved_models/checkpoint mode change 100644 => 100755 test_code/saved_models/model-33999.data-00000-of-00001 mode change 100644 => 100755 test_code/saved_models/model-33999.index mode change 100644 => 100755 test_code/test_images/actress2.jpg mode change 100644 => 100755 test_code/test_images/china6.jpg mode change 100644 => 100755 test_code/test_images/food16.jpg mode change 100644 => 100755 test_code/test_images/food6.jpg mode change 100644 => 100755 test_code/test_images/liuyifei4.jpg mode change 100644 => 100755 test_code/test_images/london1.jpg mode change 100644 => 100755 test_code/test_images/mountain4.jpg mode change 100644 => 100755 test_code/test_images/mountain5.jpg mode change 100644 => 100755 test_code/test_images/national_park1.jpg mode change 100644 => 100755 test_code/test_images/party5.jpg mode change 100644 => 100755 test_code/test_images/party7.jpg mode change 100644 => 100755 train_code/guided_filter.py mode change 100644 => 100755 train_code/layers.py mode change 100644 => 100755 train_code/loss.py mode change 100644 => 100755 train_code/network.py mode change 100644 => 100755 train_code/pretrain.py create mode 100644 train_code/selective_search/__pycache__/__init__.cpython-37.pyc create mode 100644 train_code/selective_search/__pycache__/core.cpython-37.pyc create mode 100644 train_code/selective_search/__pycache__/measure.cpython-37.pyc create mode 100644 train_code/selective_search/__pycache__/structure.cpython-37.pyc create mode 100644 train_code/selective_search/__pycache__/util.cpython-37.pyc mode change 100644 => 100755 train_code/train.py mode change 100644 => 100755 train_code/utils.py diff --git a/test_code/cartoonize.py b/test_code/cartoonize.py index f32e962..4a433c4 100755 --- a/test_code/cartoonize.py +++ b/test_code/cartoonize.py @@ -23,19 +23,19 @@ def resize_crop(image): def cartoonize(load_folder, save_folder, model_path): - input_photo = tf.placeholder(tf.float32, [1, None, None, 3]) + input_photo = tf.compat.v1.placeholder(tf.float32, [1, None, None, 3]) network_out = network.unet_generator(input_photo) final_out = guided_filter.guided_filter(input_photo, network_out, r=1, eps=5e-3) - all_vars = tf.trainable_variables() + all_vars = tf.compat.v1.trainable_variables() gene_vars = [var for var in all_vars if 'generator' in var.name] - saver = tf.train.Saver(var_list=gene_vars) + saver = tf.compat.v1.train.Saver(var_list=gene_vars) - config = tf.ConfigProto() + config = tf.compat.v1.ConfigProto() config.gpu_options.allow_growth = True - sess = tf.Session(config=config) + sess = tf.compat.v1.Session(config=config) - sess.run(tf.global_variables_initializer()) + sess.run(tf.compat.v1.global_variables_initializer()) saver.restore(sess, tf.train.latest_checkpoint(model_path)) name_list = os.listdir(load_folder) for name in tqdm(name_list): @@ -65,4 +65,3 @@ def cartoonize(load_folder, save_folder, model_path): cartoonize(load_folder, save_folder, model_path) - \ No newline at end of file diff --git a/test_code/guided_filter.py b/test_code/guided_filter.py index fd019d1..0ef511c 100755 --- a/test_code/guided_filter.py +++ b/test_code/guided_filter.py @@ -10,14 +10,14 @@ def tf_box_filter(x, r): weight = 1/(k_size**2) box_kernel = weight*np.ones((k_size, k_size, ch, 1)) box_kernel = np.array(box_kernel).astype(np.float32) - output = tf.nn.depthwise_conv2d(x, box_kernel, [1, 1, 1, 1], 'SAME') + output = tf.nn.depthwise_conv2d(input=x, filter=box_kernel, strides=[1, 1, 1, 1], padding='SAME') return output def guided_filter(x, y, r, eps=1e-2): - x_shape = tf.shape(x) + x_shape = tf.shape(input=x) #y_shape = tf.shape(y) N = tf_box_filter(tf.ones((1, x_shape[1], x_shape[2], 1), dtype=x.dtype), r) @@ -43,9 +43,9 @@ def fast_guided_filter(lr_x, lr_y, hr_x, r=1, eps=1e-8): #assert lr_x.shape.ndims == 4 and lr_y.shape.ndims == 4 and hr_x.shape.ndims == 4 - lr_x_shape = tf.shape(lr_x) + lr_x_shape = tf.shape(input=lr_x) #lr_y_shape = tf.shape(lr_y) - hr_x_shape = tf.shape(hr_x) + hr_x_shape = tf.shape(input=hr_x) N = tf_box_filter(tf.ones((1, lr_x_shape[1], lr_x_shape[2], 1), dtype=lr_x.dtype), r) @@ -57,8 +57,8 @@ def fast_guided_filter(lr_x, lr_y, hr_x, r=1, eps=1e-8): A = cov_xy / (var_x + eps) b = mean_y - A * mean_x - mean_A = tf.image.resize_images(A, hr_x_shape[1: 3]) - mean_b = tf.image.resize_images(b, hr_x_shape[1: 3]) + mean_A = tf.image.resize(A, hr_x_shape[1: 3]) + mean_b = tf.image.resize(b, hr_x_shape[1: 3]) output = mean_A * hr_x + mean_b @@ -69,17 +69,17 @@ def fast_guided_filter(lr_x, lr_y, hr_x, r=1, eps=1e-8): import cv2 from tqdm import tqdm - input_photo = tf.placeholder(tf.float32, [1, None, None, 3]) + input_photo = tf.compat.v1.placeholder(tf.float32, [1, None, None, 3]) #input_superpixel = tf.placeholder(tf.float32, [16, 256, 256, 3]) output = guided_filter(input_photo, input_photo, 5, eps=1) image = cv2.imread('output_figure1/cartoon2.jpg') image = image/127.5 - 1 image = np.expand_dims(image, axis=0) - config = tf.ConfigProto() + config = tf.compat.v1.ConfigProto() config.gpu_options.allow_growth = True - sess = tf.Session(config=config) - sess.run(tf.global_variables_initializer()) + sess = tf.compat.v1.Session(config=config) + sess.run(tf.compat.v1.global_variables_initializer()) out = sess.run(output, feed_dict={input_photo: image}) out = (np.squeeze(out)+1)*127.5 diff --git a/test_code/network.py b/test_code/network.py index 6f16cee..9aad70d 100755 --- a/test_code/network.py +++ b/test_code/network.py @@ -1,12 +1,12 @@ import tensorflow as tf import numpy as np -import tensorflow.contrib.slim as slim +import tf_slim as slim def resblock(inputs, out_channel=32, name='resblock'): - with tf.variable_scope(name): + with tf.compat.v1.variable_scope(name): x = slim.convolution2d(inputs, out_channel, [3, 3], activation_fn=None, scope='conv1') @@ -20,7 +20,7 @@ def resblock(inputs, out_channel=32, name='resblock'): def unet_generator(inputs, channel=32, num_blocks=4, name='generator', reuse=False): - with tf.variable_scope(name, reuse=reuse): + with tf.compat.v1.variable_scope(name, reuse=reuse): x0 = slim.convolution2d(inputs, channel, [7, 7], activation_fn=None) x0 = tf.nn.leaky_relu(x0) @@ -41,15 +41,15 @@ def unet_generator(inputs, channel=32, num_blocks=4, name='generator', reuse=Fal x2 = slim.convolution2d(x2, channel*2, [3, 3], activation_fn=None) x2 = tf.nn.leaky_relu(x2) - h1, w1 = tf.shape(x2)[1], tf.shape(x2)[2] - x3 = tf.image.resize_bilinear(x2, (h1*2, w1*2)) + h1, w1 = tf.shape(input=x2)[1], tf.shape(input=x2)[2] + x3 = tf.image.resize(x2, (h1*2, w1*2), method=tf.image.ResizeMethod.BILINEAR) x3 = slim.convolution2d(x3+x1, channel*2, [3, 3], activation_fn=None) x3 = tf.nn.leaky_relu(x3) x3 = slim.convolution2d(x3, channel, [3, 3], activation_fn=None) x3 = tf.nn.leaky_relu(x3) - h2, w2 = tf.shape(x3)[1], tf.shape(x3)[2] - x4 = tf.image.resize_bilinear(x3, (h2*2, w2*2)) + h2, w2 = tf.shape(input=x3)[1], tf.shape(input=x3)[2] + x4 = tf.image.resize(x3, (h2*2, w2*2), method=tf.image.ResizeMethod.BILINEAR) x4 = slim.convolution2d(x4+x0, channel, [3, 3], activation_fn=None) x4 = tf.nn.leaky_relu(x4) x4 = slim.convolution2d(x4, 3, [7, 7], activation_fn=None) diff --git a/test_code/saved_models/checkpoint b/test_code/saved_models/checkpoint old mode 100644 new mode 100755 diff --git a/test_code/saved_models/model-33999.data-00000-of-00001 b/test_code/saved_models/model-33999.data-00000-of-00001 old mode 100644 new mode 100755 diff --git a/test_code/saved_models/model-33999.index b/test_code/saved_models/model-33999.index old mode 100644 new mode 100755 diff --git a/test_code/test_images/actress2.jpg b/test_code/test_images/actress2.jpg old mode 100644 new mode 100755 diff --git a/test_code/test_images/china6.jpg b/test_code/test_images/china6.jpg old mode 100644 new mode 100755 diff --git a/test_code/test_images/food16.jpg b/test_code/test_images/food16.jpg old mode 100644 new mode 100755 diff --git a/test_code/test_images/food6.jpg b/test_code/test_images/food6.jpg old mode 100644 new mode 100755 diff --git a/test_code/test_images/liuyifei4.jpg b/test_code/test_images/liuyifei4.jpg old mode 100644 new mode 100755 diff --git a/test_code/test_images/london1.jpg b/test_code/test_images/london1.jpg old mode 100644 new mode 100755 diff --git a/test_code/test_images/mountain4.jpg b/test_code/test_images/mountain4.jpg old mode 100644 new mode 100755 diff --git a/test_code/test_images/mountain5.jpg b/test_code/test_images/mountain5.jpg old mode 100644 new mode 100755 diff --git a/test_code/test_images/national_park1.jpg b/test_code/test_images/national_park1.jpg old mode 100644 new mode 100755 diff --git a/test_code/test_images/party5.jpg b/test_code/test_images/party5.jpg old mode 100644 new mode 100755 diff --git a/test_code/test_images/party7.jpg b/test_code/test_images/party7.jpg old mode 100644 new mode 100755 diff --git a/train_code/guided_filter.py b/train_code/guided_filter.py old mode 100644 new mode 100755 index 340f450..5f86d86 --- a/train_code/guided_filter.py +++ b/train_code/guided_filter.py @@ -13,14 +13,14 @@ def tf_box_filter(x, r): weight = 1/((2*r+1)**2) box_kernel = weight*np.ones((2*r+1, 2*r+1, ch, 1)) box_kernel = np.array(box_kernel).astype(np.float32) - output = tf.nn.depthwise_conv2d(x, box_kernel, [1, 1, 1, 1], 'SAME') + output = tf.nn.depthwise_conv2d(input=x, filter=box_kernel, strides=[1, 1, 1, 1], padding='SAME') return output def guided_filter(x, y, r, eps=1e-2): - x_shape = tf.shape(x) + x_shape = tf.shape(input=x) #y_shape = tf.shape(y) N = tf_box_filter(tf.ones((1, x_shape[1], x_shape[2], 1), dtype=x.dtype), r) diff --git a/train_code/layers.py b/train_code/layers.py old mode 100644 new mode 100755 index e361ad5..cf1a854 --- a/train_code/layers.py +++ b/train_code/layers.py @@ -1,93 +1,92 @@ -''' -CVPR 2020 submission, Paper ID 6791 -Source code for 'Learning to Cartoonize Using White-Box Cartoon Representations' -''' - - -import tensorflow as tf -import numpy as np -import tensorflow.contrib.slim as slim - - - -def adaptive_instance_norm(content, style, epsilon=1e-5): - - c_mean, c_var = tf.nn.moments(content, axes=[1, 2], keep_dims=True) - s_mean, s_var = tf.nn.moments(style, axes=[1, 2], keep_dims=True) - c_std, s_std = tf.sqrt(c_var + epsilon), tf.sqrt(s_var + epsilon) - - return s_std * (content - c_mean) / c_std + s_mean - - - -def spectral_norm(w, iteration=1): - w_shape = w.shape.as_list() - w = tf.reshape(w, [-1, w_shape[-1]]) - - u = tf.get_variable("u", [1, w_shape[-1]], - initializer=tf.random_normal_initializer(), trainable=False) - - u_hat = u - v_hat = None - for i in range(iteration): - """ - power iteration - Usually iteration = 1 will be enough - """ - v_ = tf.matmul(u_hat, tf.transpose(w)) - v_hat = tf.nn.l2_normalize(v_) - - u_ = tf.matmul(v_hat, w) - u_hat = tf.nn.l2_normalize(u_) - - u_hat = tf.stop_gradient(u_hat) - v_hat = tf.stop_gradient(v_hat) - - sigma = tf.matmul(tf.matmul(v_hat, w), tf.transpose(u_hat)) - - with tf.control_dependencies([u.assign(u_hat)]): - w_norm = w / sigma - w_norm = tf.reshape(w_norm, w_shape) - - return w_norm - - -def conv_spectral_norm(x, channel, k_size, stride=1, name='conv_snorm'): - with tf.variable_scope(name): - w = tf.get_variable("kernel", shape=[k_size[0], k_size[1], x.get_shape()[-1], channel]) - b = tf.get_variable("bias", [channel], initializer=tf.constant_initializer(0.0)) - - x = tf.nn.conv2d(input=x, filter=spectral_norm(w), strides=[1, stride, stride, 1], padding='SAME') + b - - return x - - - -def self_attention(inputs, name='attention', reuse=False): - with tf.variable_scope(name, reuse=reuse): - h, w = tf.shape(inputs)[1], tf.shape(inputs)[2] - bs, _, _, ch = inputs.get_shape().as_list() - f = slim.convolution2d(inputs, ch//8, [1, 1], activation_fn=None) - g = slim.convolution2d(inputs, ch//8, [1, 1], activation_fn=None) - s = slim.convolution2d(inputs, 1, [1, 1], activation_fn=None) - f_flatten = tf.reshape(f, shape=[f.shape[0], -1, f.shape[-1]]) - g_flatten = tf.reshape(g, shape=[g.shape[0], -1, g.shape[-1]]) - beta = tf.matmul(f_flatten, g_flatten, transpose_b=True) - beta = tf.nn.softmax(beta) - - s_flatten = tf.reshape(s, shape=[s.shape[0], -1, s.shape[-1]]) - att_map = tf.matmul(beta, s_flatten) - att_map = tf.reshape(att_map, shape=[bs, h, w, 1]) - gamma = tf.get_variable("gamma", [1], initializer=tf.constant_initializer(0.0)) - output = att_map * gamma + inputs - - return att_map, output - - - -if __name__ == '__main__': - pass - - - - \ No newline at end of file +''' +CVPR 2020 submission, Paper ID 6791 +Source code for 'Learning to Cartoonize Using White-Box Cartoon Representations' +''' + + +import tensorflow as tf +import numpy as np +import tf_slim as slim + + + +def adaptive_instance_norm(content, style, epsilon=1e-5): + + c_mean, c_var = tf.nn.moments(x=content, axes=[1, 2], keepdims=True) + s_mean, s_var = tf.nn.moments(x=style, axes=[1, 2], keepdims=True) + c_std, s_std = tf.sqrt(c_var + epsilon), tf.sqrt(s_var + epsilon) + + return s_std * (content - c_mean) / c_std + s_mean + + + +def spectral_norm(w, iteration=1): + w_shape = w.shape.as_list() + w = tf.reshape(w, [-1, w_shape[-1]]) + + u = tf.compat.v1.get_variable("u", [1, w_shape[-1]], + initializer=tf.compat.v1.random_normal_initializer(), trainable=False) + + u_hat = u + v_hat = None + for i in range(iteration): + """ + power iteration + Usually iteration = 1 will be enough + """ + v_ = tf.matmul(u_hat, tf.transpose(a=w)) + v_hat = tf.nn.l2_normalize(v_) + + u_ = tf.matmul(v_hat, w) + u_hat = tf.nn.l2_normalize(u_) + + u_hat = tf.stop_gradient(u_hat) + v_hat = tf.stop_gradient(v_hat) + + sigma = tf.matmul(tf.matmul(v_hat, w), tf.transpose(a=u_hat)) + + with tf.control_dependencies([u.assign(u_hat)]): + w_norm = w / sigma + w_norm = tf.reshape(w_norm, w_shape) + + return w_norm + + +def conv_spectral_norm(x, channel, k_size, stride=1, name='conv_snorm'): + with tf.compat.v1.variable_scope(name): + w = tf.compat.v1.get_variable("kernel", shape=[k_size[0], k_size[1], x.get_shape()[-1], channel]) + b = tf.compat.v1.get_variable("bias", [channel], initializer=tf.compat.v1.constant_initializer(0.0)) + + x = tf.nn.conv2d(input=x, filters=spectral_norm(w), strides=[1, stride, stride, 1], padding='SAME') + b + + return x + + + +def self_attention(inputs, name='attention', reuse=False): + with tf.compat.v1.variable_scope(name, reuse=reuse): + h, w = tf.shape(input=inputs)[1], tf.shape(input=inputs)[2] + bs, _, _, ch = inputs.get_shape().as_list() + f = slim.convolution2d(inputs, ch//8, [1, 1], activation_fn=None) + g = slim.convolution2d(inputs, ch//8, [1, 1], activation_fn=None) + s = slim.convolution2d(inputs, 1, [1, 1], activation_fn=None) + f_flatten = tf.reshape(f, shape=[f.shape[0], -1, f.shape[-1]]) + g_flatten = tf.reshape(g, shape=[g.shape[0], -1, g.shape[-1]]) + beta = tf.matmul(f_flatten, g_flatten, transpose_b=True) + beta = tf.nn.softmax(beta) + + s_flatten = tf.reshape(s, shape=[s.shape[0], -1, s.shape[-1]]) + att_map = tf.matmul(beta, s_flatten) + att_map = tf.reshape(att_map, shape=[bs, h, w, 1]) + gamma = tf.compat.v1.get_variable("gamma", [1], initializer=tf.compat.v1.constant_initializer(0.0)) + output = att_map * gamma + inputs + + return att_map, output + + + +if __name__ == '__main__': + pass + + + diff --git a/train_code/loss.py b/train_code/loss.py old mode 100644 new mode 100755 index 889a243..6e95c43 --- a/train_code/loss.py +++ b/train_code/loss.py @@ -1,191 +1,190 @@ -''' -CVPR 2020 submission, Paper ID 6791 -Source code for 'Learning to Cartoonize Using White-Box Cartoon Representations' -''' - - -import numpy as np -import scipy.stats as st -import tensorflow as tf - - - -VGG_MEAN = [103.939, 116.779, 123.68] - - -class Vgg19: - - def __init__(self, vgg19_npy_path=None): - - self.data_dict = np.load(vgg19_npy_path, encoding='latin1', allow_pickle=True).item() - print('Finished loading vgg19.npy') - - - def build_conv4_4(self, rgb, include_fc=False): - - rgb_scaled = (rgb+1) * 127.5 - - blue, green, red = tf.split(axis=3, num_or_size_splits=3, value=rgb_scaled) - bgr = tf.concat(axis=3, values=[blue - VGG_MEAN[0], - green - VGG_MEAN[1], red - VGG_MEAN[2]]) - - self.conv1_1 = self.conv_layer(bgr, "conv1_1") - self.relu1_1 = tf.nn.relu(self.conv1_1) - self.conv1_2 = self.conv_layer(self.relu1_1, "conv1_2") - self.relu1_2 = tf.nn.relu(self.conv1_2) - self.pool1 = self.max_pool(self.relu1_2, 'pool1') - - self.conv2_1 = self.conv_layer(self.pool1, "conv2_1") - self.relu2_1 = tf.nn.relu(self.conv2_1) - self.conv2_2 = self.conv_layer(self.relu2_1, "conv2_2") - self.relu2_2 = tf.nn.relu(self.conv2_2) - self.pool2 = self.max_pool(self.relu2_2, 'pool2') - - self.conv3_1 = self.conv_layer(self.pool2, "conv3_1") - self.relu3_1 = tf.nn.relu(self.conv3_1) - self.conv3_2 = self.conv_layer(self.relu3_1, "conv3_2") - self.relu3_2 = tf.nn.relu(self.conv3_2) - self.conv3_3 = self.conv_layer(self.relu3_2, "conv3_3") - self.relu3_3 = tf.nn.relu(self.conv3_3) - self.conv3_4 = self.conv_layer(self.relu3_3, "conv3_4") - self.relu3_4 = tf.nn.relu(self.conv3_4) - self.pool3 = self.max_pool(self.relu3_4, 'pool3') - - self.conv4_1 = self.conv_layer(self.pool3, "conv4_1") - self.relu4_1 = tf.nn.relu(self.conv4_1) - self.conv4_2 = self.conv_layer(self.relu4_1, "conv4_2") - self.relu4_2 = tf.nn.relu(self.conv4_2) - self.conv4_3 = self.conv_layer(self.relu4_2, "conv4_3") - self.relu4_3 = tf.nn.relu(self.conv4_3) - self.conv4_4 = self.conv_layer(self.relu4_3, "conv4_4") - self.relu4_4 = tf.nn.relu(self.conv4_4) - self.pool4 = self.max_pool(self.relu4_4, 'pool4') - - return self.conv4_4 - - def max_pool(self, bottom, name): - return tf.nn.max_pool(bottom, ksize=[1, 2, 2, 1], - strides=[1, 2, 2, 1], padding='SAME', name=name) - - def conv_layer(self, bottom, name): - with tf.variable_scope(name): - filt = self.get_conv_filter(name) - - conv = tf.nn.conv2d(bottom, filt, [1, 1, 1, 1], padding='SAME') - - conv_biases = self.get_bias(name) - bias = tf.nn.bias_add(conv, conv_biases) - - #relu = tf.nn.relu(bias) - return bias - - - - def fc_layer(self, bottom, name): - with tf.variable_scope(name): - shape = bottom.get_shape().as_list() - dim = 1 - for d in shape[1:]: - dim *= d - x = tf.reshape(bottom, [-1, dim]) - - weights = self.get_fc_weight(name) - biases = self.get_bias(name) - - # Fully connected layer. Note that the '+' operation automatically - # broadcasts the biases. - fc = tf.nn.bias_add(tf.matmul(x, weights), biases) - - return fc - - def get_conv_filter(self, name): - return tf.constant(self.data_dict[name][0], name="filter") - - def get_bias(self, name): - return tf.constant(self.data_dict[name][1], name="biases") - - def get_fc_weight(self, name): - return tf.constant(self.data_dict[name][0], name="weights") - - - -def vggloss_4_4(image_a, image_b): - vgg_model = Vgg19('vgg19_no_fc.npy') - vgg_a = vgg_model.build_conv4_4(image_a) - vgg_b = vgg_model.build_conv4_4(image_b) - VGG_loss = tf.losses.absolute_difference(vgg_a, vgg_b) - #VGG_loss = tf.nn.l2_loss(vgg_a - vgg_b) - h, w, c= vgg_a.get_shape().as_list()[1:] - VGG_loss = tf.reduce_mean(VGG_loss)/(h*w*c) - return VGG_loss - - - -def wgan_loss(discriminator, real, fake, patch=True, - channel=32, name='discriminator', lambda_=2): - real_logits = discriminator(real, patch=patch, channel=channel, name=name, reuse=False) - fake_logits = discriminator(fake, patch=patch, channel=channel, name=name, reuse=True) - - d_loss_real = - tf.reduce_mean(real_logits) - d_loss_fake = tf.reduce_mean(fake_logits) - - d_loss = d_loss_real + d_loss_fake - g_loss = - d_loss_fake - - """ Gradient Penalty """ - # This is borrowed from https://github.com/kodalinaveen3/DRAGAN/blob/master/DRAGAN.ipynb - alpha = tf.random_uniform([tf.shape(real)[0], 1, 1, 1], minval=0.,maxval=1.) - differences = fake - real # This is different from MAGAN - interpolates = real + (alpha * differences) - inter_logit = discriminator(interpolates, channel=channel, name=name, reuse=True) - gradients = tf.gradients(inter_logit, [interpolates])[0] - slopes = tf.sqrt(tf.reduce_sum(tf.square(gradients), reduction_indices=[1])) - gradient_penalty = tf.reduce_mean((slopes - 1.) ** 2) - d_loss += lambda_ * gradient_penalty - - return d_loss, g_loss - - -def gan_loss(discriminator, real, fake, scale=1,channel=32, patch=False, name='discriminator'): - - real_logit = discriminator(real, scale, channel, name=name, patch=patch, reuse=False) - fake_logit = discriminator(fake, scale, channel, name=name, patch=patch, reuse=True) - - real_logit = tf.nn.sigmoid(real_logit) - fake_logit = tf.nn.sigmoid(fake_logit) - - g_loss_blur = -tf.reduce_mean(tf.log(fake_logit)) - d_loss_blur = -tf.reduce_mean(tf.log(real_logit) + tf.log(1. - fake_logit)) - - return d_loss_blur, g_loss_blur - - - -def lsgan_loss(discriminator, real, fake, scale=1, - channel=32, patch=False, name='discriminator'): - - real_logit = discriminator(real, scale, channel, name=name, patch=patch, reuse=False) - fake_logit = discriminator(fake, scale, channel, name=name, patch=patch, reuse=True) - - g_loss = tf.reduce_mean((fake_logit - 1)**2) - d_loss = 0.5*(tf.reduce_mean((real_logit - 1)**2) + tf.reduce_mean(fake_logit**2)) - - return d_loss, g_loss - - - -def total_variation_loss(image, k_size=1): - h, w = image.get_shape().as_list()[1:3] - tv_h = tf.reduce_mean((image[:, k_size:, :, :] - image[:, :h - k_size, :, :])**2) - tv_w = tf.reduce_mean((image[:, :, k_size:, :] - image[:, :, :w - k_size, :])**2) - tv_loss = (tv_h + tv_w)/(3*h*w) - return tv_loss - - - - -if __name__ == '__main__': - pass - - - \ No newline at end of file +''' +CVPR 2020 submission, Paper ID 6791 +Source code for 'Learning to Cartoonize Using White-Box Cartoon Representations' +''' + + +import numpy as np +import scipy.stats as st +import tensorflow as tf + + + +VGG_MEAN = [103.939, 116.779, 123.68] + + +class Vgg19: + + def __init__(self, vgg19_npy_path=None): + + self.data_dict = np.load(vgg19_npy_path, encoding='latin1', allow_pickle=True).item() + print('Finished loading vgg19.npy') + + + def build_conv4_4(self, rgb, include_fc=False): + + rgb_scaled = (rgb+1) * 127.5 + + blue, green, red = tf.split(axis=3, num_or_size_splits=3, value=rgb_scaled) + bgr = tf.concat(axis=3, values=[blue - VGG_MEAN[0], + green - VGG_MEAN[1], red - VGG_MEAN[2]]) + + self.conv1_1 = self.conv_layer(bgr, "conv1_1") + self.relu1_1 = tf.nn.relu(self.conv1_1) + self.conv1_2 = self.conv_layer(self.relu1_1, "conv1_2") + self.relu1_2 = tf.nn.relu(self.conv1_2) + self.pool1 = self.max_pool(self.relu1_2, 'pool1') + + self.conv2_1 = self.conv_layer(self.pool1, "conv2_1") + self.relu2_1 = tf.nn.relu(self.conv2_1) + self.conv2_2 = self.conv_layer(self.relu2_1, "conv2_2") + self.relu2_2 = tf.nn.relu(self.conv2_2) + self.pool2 = self.max_pool(self.relu2_2, 'pool2') + + self.conv3_1 = self.conv_layer(self.pool2, "conv3_1") + self.relu3_1 = tf.nn.relu(self.conv3_1) + self.conv3_2 = self.conv_layer(self.relu3_1, "conv3_2") + self.relu3_2 = tf.nn.relu(self.conv3_2) + self.conv3_3 = self.conv_layer(self.relu3_2, "conv3_3") + self.relu3_3 = tf.nn.relu(self.conv3_3) + self.conv3_4 = self.conv_layer(self.relu3_3, "conv3_4") + self.relu3_4 = tf.nn.relu(self.conv3_4) + self.pool3 = self.max_pool(self.relu3_4, 'pool3') + + self.conv4_1 = self.conv_layer(self.pool3, "conv4_1") + self.relu4_1 = tf.nn.relu(self.conv4_1) + self.conv4_2 = self.conv_layer(self.relu4_1, "conv4_2") + self.relu4_2 = tf.nn.relu(self.conv4_2) + self.conv4_3 = self.conv_layer(self.relu4_2, "conv4_3") + self.relu4_3 = tf.nn.relu(self.conv4_3) + self.conv4_4 = self.conv_layer(self.relu4_3, "conv4_4") + self.relu4_4 = tf.nn.relu(self.conv4_4) + self.pool4 = self.max_pool(self.relu4_4, 'pool4') + + return self.conv4_4 + + def max_pool(self, bottom, name): + return tf.nn.max_pool2d(input=bottom, ksize=[1, 2, 2, 1], + strides=[1, 2, 2, 1], padding='SAME', name=name) + + def conv_layer(self, bottom, name): + with tf.compat.v1.variable_scope(name): + filt = self.get_conv_filter(name) + + conv = tf.nn.conv2d(input=bottom, filters=filt, strides=[1, 1, 1, 1], padding='SAME') + + conv_biases = self.get_bias(name) + bias = tf.nn.bias_add(conv, conv_biases) + + #relu = tf.nn.relu(bias) + return bias + + + + def fc_layer(self, bottom, name): + with tf.compat.v1.variable_scope(name): + shape = bottom.get_shape().as_list() + dim = 1 + for d in shape[1:]: + dim *= d + x = tf.reshape(bottom, [-1, dim]) + + weights = self.get_fc_weight(name) + biases = self.get_bias(name) + + # Fully connected layer. Note that the '+' operation automatically + # broadcasts the biases. + fc = tf.nn.bias_add(tf.matmul(x, weights), biases) + + return fc + + def get_conv_filter(self, name): + return tf.constant(self.data_dict[name][0], name="filter") + + def get_bias(self, name): + return tf.constant(self.data_dict[name][1], name="biases") + + def get_fc_weight(self, name): + return tf.constant(self.data_dict[name][0], name="weights") + + + +def vggloss_4_4(image_a, image_b): + vgg_model = Vgg19('vgg19_no_fc.npy') + vgg_a = vgg_model.build_conv4_4(image_a) + vgg_b = vgg_model.build_conv4_4(image_b) + VGG_loss = tf.compat.v1.losses.absolute_difference(vgg_a, vgg_b) + #VGG_loss = tf.nn.l2_loss(vgg_a - vgg_b) + h, w, c= vgg_a.get_shape().as_list()[1:] + VGG_loss = tf.reduce_mean(input_tensor=VGG_loss)/(h*w*c) + return VGG_loss + + + +def wgan_loss(discriminator, real, fake, patch=True, + channel=32, name='discriminator', lambda_=2): + real_logits = discriminator(real, patch=patch, channel=channel, name=name, reuse=False) + fake_logits = discriminator(fake, patch=patch, channel=channel, name=name, reuse=True) + + d_loss_real = - tf.reduce_mean(input_tensor=real_logits) + d_loss_fake = tf.reduce_mean(input_tensor=fake_logits) + + d_loss = d_loss_real + d_loss_fake + g_loss = - d_loss_fake + + """ Gradient Penalty """ + # This is borrowed from https://github.com/kodalinaveen3/DRAGAN/blob/master/DRAGAN.ipynb + alpha = tf.random.uniform([tf.shape(input=real)[0], 1, 1, 1], minval=0.,maxval=1.) + differences = fake - real # This is different from MAGAN + interpolates = real + (alpha * differences) + inter_logit = discriminator(interpolates, channel=channel, name=name, reuse=True) + gradients = tf.gradients(ys=inter_logit, xs=[interpolates])[0] + slopes = tf.sqrt(tf.reduce_sum(input_tensor=tf.square(gradients), axis=[1])) + gradient_penalty = tf.reduce_mean(input_tensor=(slopes - 1.) ** 2) + d_loss += lambda_ * gradient_penalty + + return d_loss, g_loss + + +def gan_loss(discriminator, real, fake, scale=1,channel=32, patch=False, name='discriminator'): + + real_logit = discriminator(real, scale, channel, name=name, patch=patch, reuse=False) + fake_logit = discriminator(fake, scale, channel, name=name, patch=patch, reuse=True) + + real_logit = tf.nn.sigmoid(real_logit) + fake_logit = tf.nn.sigmoid(fake_logit) + + g_loss_blur = -tf.reduce_mean(input_tensor=tf.math.log(fake_logit)) + d_loss_blur = -tf.reduce_mean(input_tensor=tf.math.log(real_logit) + tf.math.log(1. - fake_logit)) + + return d_loss_blur, g_loss_blur + + + +def lsgan_loss(discriminator, real, fake, scale=1, + channel=32, patch=False, name='discriminator'): + + real_logit = discriminator(real, scale, channel, name=name, patch=patch, reuse=False) + fake_logit = discriminator(fake, scale, channel, name=name, patch=patch, reuse=True) + + g_loss = tf.reduce_mean(input_tensor=(fake_logit - 1)**2) + d_loss = 0.5*(tf.reduce_mean(input_tensor=(real_logit - 1)**2) + tf.reduce_mean(input_tensor=fake_logit**2)) + + return d_loss, g_loss + + + +def total_variation_loss(image, k_size=1): + h, w = image.get_shape().as_list()[1:3] + tv_h = tf.reduce_mean(input_tensor=(image[:, k_size:, :, :] - image[:, :h - k_size, :, :])**2) + tv_w = tf.reduce_mean(input_tensor=(image[:, :, k_size:, :] - image[:, :, :w - k_size, :])**2) + tv_loss = (tv_h + tv_w)/(3*h*w) + return tv_loss + + + + +if __name__ == '__main__': + pass + + diff --git a/train_code/network.py b/train_code/network.py old mode 100644 new mode 100755 index 8527aa8..dc10944 --- a/train_code/network.py +++ b/train_code/network.py @@ -7,7 +7,7 @@ import layers import tensorflow as tf import numpy as np -import tensorflow.contrib.slim as slim +import tf_slim as slim from tqdm import tqdm @@ -15,7 +15,7 @@ def resblock(inputs, out_channel=32, name='resblock'): - with tf.variable_scope(name): + with tf.compat.v1.variable_scope(name): x = slim.convolution2d(inputs, out_channel, [3, 3], activation_fn=None, scope='conv1') @@ -28,7 +28,7 @@ def resblock(inputs, out_channel=32, name='resblock'): def generator(inputs, channel=32, num_blocks=4, name='generator', reuse=False): - with tf.variable_scope(name, reuse=reuse): + with tf.compat.v1.variable_scope(name, reuse=reuse): x = slim.convolution2d(inputs, channel, [7, 7], activation_fn=None) x = tf.nn.leaky_relu(x) @@ -60,7 +60,7 @@ def generator(inputs, channel=32, num_blocks=4, name='generator', reuse=False): def unet_generator(inputs, channel=32, num_blocks=4, name='generator', reuse=False): - with tf.variable_scope(name, reuse=reuse): + with tf.compat.v1.variable_scope(name, reuse=reuse): x0 = slim.convolution2d(inputs, channel, [7, 7], activation_fn=None) x0 = tf.nn.leaky_relu(x0) @@ -81,15 +81,15 @@ def unet_generator(inputs, channel=32, num_blocks=4, name='generator', reuse=Fal x2 = slim.convolution2d(x2, channel*2, [3, 3], activation_fn=None) x2 = tf.nn.leaky_relu(x2) - h1, w1 = tf.shape(x2)[1], tf.shape(x2)[2] - x3 = tf.image.resize_bilinear(x2, (h1*2, w1*2)) + h1, w1 = tf.shape(input=x2)[1], tf.shape(input=x2)[2] + x3 = tf.image.resize(x2, (h1*2, w1*2), method=tf.image.ResizeMethod.BILINEAR) x3 = slim.convolution2d(x3+x1, channel*2, [3, 3], activation_fn=None) x3 = tf.nn.leaky_relu(x3) x3 = slim.convolution2d(x3, channel, [3, 3], activation_fn=None) x3 = tf.nn.leaky_relu(x3) - h2, w2 = tf.shape(x3)[1], tf.shape(x3)[2] - x4 = tf.image.resize_bilinear(x3, (h2*2, w2*2)) + h2, w2 = tf.shape(input=x3)[1], tf.shape(input=x3)[2] + x4 = tf.image.resize(x3, (h2*2, w2*2), method=tf.image.ResizeMethod.BILINEAR) x4 = slim.convolution2d(x4+x0, channel, [3, 3], activation_fn=None) x4 = tf.nn.leaky_relu(x4) x4 = slim.convolution2d(x4, 3, [7, 7], activation_fn=None) @@ -101,7 +101,7 @@ def unet_generator(inputs, channel=32, num_blocks=4, name='generator', reuse=Fal def disc_bn(x, scale=1, channel=32, is_training=True, name='discriminator', patch=True, reuse=False): - with tf.variable_scope(name, reuse=reuse): + with tf.compat.v1.variable_scope(name, reuse=reuse): for idx in range(3): x = slim.convolution2d(x, channel*2**idx, [3, 3], stride=2, activation_fn=None) @@ -115,7 +115,7 @@ def disc_bn(x, scale=1, channel=32, is_training=True, if patch == True: x = slim.convolution2d(x, 1, [1, 1], activation_fn=None) else: - x = tf.reduce_mean(x, axis=[1, 2]) + x = tf.reduce_mean(input_tensor=x, axis=[1, 2]) x = slim.fully_connected(x, 1, activation_fn=None) return x @@ -124,7 +124,7 @@ def disc_bn(x, scale=1, channel=32, is_training=True, def disc_sn(x, scale=1, channel=32, patch=True, name='discriminator', reuse=False): - with tf.variable_scope(name, reuse=reuse): + with tf.compat.v1.variable_scope(name, reuse=reuse): for idx in range(3): x = layers.conv_spectral_norm(x, channel*2**idx, [3, 3], @@ -140,14 +140,14 @@ def disc_sn(x, scale=1, channel=32, patch=True, name='discriminator', reuse=Fals x = layers.conv_spectral_norm(x, 1, [1, 1], name='conv_out'.format(idx)) else: - x = tf.reduce_mean(x, axis=[1, 2]) + x = tf.reduce_mean(input_tensor=x, axis=[1, 2]) x = slim.fully_connected(x, 1, activation_fn=None) return x def disc_ln(x, channel=32, is_training=True, name='discriminator', patch=True, reuse=False): - with tf.variable_scope(name, reuse=reuse): + with tf.compat.v1.variable_scope(name, reuse=reuse): for idx in range(3): x = slim.convolution2d(x, channel*2**idx, [3, 3], stride=2, activation_fn=None) @@ -161,7 +161,7 @@ def disc_ln(x, channel=32, is_training=True, name='discriminator', patch=True, r if patch == True: x = slim.convolution2d(x, 1, [1, 1], activation_fn=None) else: - x = tf.reduce_mean(x, axis=[1, 2]) + x = tf.reduce_mean(input_tensor=x, axis=[1, 2]) x = slim.fully_connected(x, 1, activation_fn=None) return x @@ -172,4 +172,3 @@ def disc_ln(x, channel=32, is_training=True, name='discriminator', patch=True, r if __name__ == '__main__': pass - \ No newline at end of file diff --git a/train_code/pretrain.py b/train_code/pretrain.py old mode 100644 new mode 100755 index d18684b..0f694af --- a/train_code/pretrain.py +++ b/train_code/pretrain.py @@ -7,8 +7,7 @@ import tensorflow as tf -import tensorflow.contrib.slim as slim - +import tf_slim as slim import utils import os import numpy as np @@ -19,7 +18,6 @@ os.environ["CUDA_VISIBLE_DEVICES"]="0" - def arg_parser(): parser = argparse.ArgumentParser() parser.add_argument("--patch_size", default = 256, type = int) @@ -38,20 +36,20 @@ def arg_parser(): def train(args): - input_photo = tf.placeholder(tf.float32, [args.batch_size, + input_photo = tf.compat.v1.placeholder(tf.float32, [args.batch_size, args.patch_size, args.patch_size, 3]) output = network.unet_generator(input_photo) - recon_loss = tf.reduce_mean(tf.losses.absolute_difference(input_photo, output)) + recon_loss = tf.reduce_mean(input_tensor=tf.compat.v1.losses.absolute_difference(input_photo, output)) - all_vars = tf.trainable_variables() + all_vars = tf.compat.v1.trainable_variables() gene_vars = [var for var in all_vars if 'gene' in var.name] - update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS) + update_ops = tf.compat.v1.get_collection(tf.compat.v1.GraphKeys.UPDATE_OPS) with tf.control_dependencies(update_ops): - optim = tf.train.AdamOptimizer(args.adv_train_lr, beta1=0.5, beta2=0.99)\ + optim = tf.compat.v1.train.AdamOptimizer(args.adv_train_lr, beta1=0.5, beta2=0.99)\ .minimize(recon_loss, var_list=gene_vars) @@ -60,17 +58,17 @@ def train(args): config.gpu_options.allow_growth = True sess = tf.Session(config=config) ''' - gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=args.gpu_fraction) - sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options)) - saver = tf.train.Saver(var_list=gene_vars, max_to_keep=20) + gpu_options = tf.compat.v1.GPUOptions(per_process_gpu_memory_fraction=args.gpu_fraction) + sess = tf.compat.v1.Session(config=tf.compat.v1.ConfigProto(gpu_options=gpu_options)) + saver = tf.compat.v1.train.Saver(var_list=gene_vars, max_to_keep=20) with tf.device('/device:GPU:0'): - sess.run(tf.global_variables_initializer()) - face_photo_dir = 'dataset/photo_face' - face_photo_list = utils.load_image_list(face_photo_dir) - scenery_photo_dir = 'dataset/photo_scenery' - scenery_photo_list = utils.load_image_list(scenery_photo_dir) + sess.run(tf.compat.v1.global_variables_initializer()) + face_photo_dir = 'dataset/face_photo' + utils.load_image_list(face_photo_dir) + scenery_photo_dir = 'dataset/scenery_photo' + utils.load_image_list(scenery_photo_dir) for total_iter in tqdm(range(args.total_iter)): @@ -84,8 +82,10 @@ def train(args): if np.mod(total_iter+1, 50) == 0: + wandb.log({"r_loss":r_loss,"iteration":total_iter}) print('pretrain, iter: {}, recon_loss: {}'.format(total_iter, r_loss)) if np.mod(total_iter+1, 500 ) == 0: + saver.save(sess, args.save_dir+'save_models/model', write_meta_graph=False, global_step=total_iter) @@ -104,6 +104,7 @@ def train(args): str(total_iter)+'_scenery_result.jpg', 4) utils.write_batch_image(photo_scenery, args.save_dir+'/images', str(total_iter)+'_scenery_photo.jpg', 4) + @@ -113,4 +114,3 @@ def train(args): args = arg_parser() train(args) - \ No newline at end of file diff --git a/train_code/selective_search/__pycache__/__init__.cpython-37.pyc b/train_code/selective_search/__pycache__/__init__.cpython-37.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fc2cd4a98a70ead66efe72ea9692dac9d6cb0ada GIT binary patch literal 267 zcmZ8cu}TCn5KXcQa&T~kKe3Gi)*^==a4cSTO^BJ14Q4mXBzJnBc zd91o;MJ_RWq6`}W0>nTLIf|D6x#y7I((o-O{smoh(q5ANN|mEwTJ(D!smDXU@bTWmtN#H~Iki)>Q$BGj zh89Q5t2kO-?Vh?5zxJQGXAX1l_2_&2Hd_9fqX+l~XU=^-2__A+8+A}O&fqt+!wc;u z&3b(6V|Z8ar{F${o~@iY?>eUpUWH#jaj!a)uxeGS&(`YjjGZ#uyCVE*y$&JezwVrJ zj5h1AZXR%WGf1{L_R_|)Co|cr(kBuv+I4%RvEMe4Otu9WN~Z_#1$nSN_miDUk;r8Q z0o(oI;;~&?Jk8@$rTzBL+J#7zO2!Kgao-0kEs_jkNnb8l7ovlS&s5}nnn*HoEFuWBtIM|sRv~y+QYf6?SoQb(L`0WO(px-)^ZUIi@~H+_$uA`K5Qmu z=yrJkKjmqbsI(flbD8K7^hMT$4`ek`vz2F#S(vt89!rh&vs8cC-n-?f5M0_#@; z`I3=CmBOpl_>u+oa@0n9F3>JU`My+ei=Y5P=OFD5Ffc52YU7y}R>NSH^3(YAsDS7} zjyRO~dwBIGNFT}mpfzDN6V6rVbN<*vGBM-A+`m(C@|!<0>{zWp)RO1L_Mpt=woZi@ zq}xZOI9q_%0dn+!38ZX&k#s&OW7?Iwt(-tLtJQm~R@8zq_+s9k~e z=1Ch*XG2QKo5S(*=r;MZSLVYzKP97S2Z3{Pi_2?8d_+RoaB9j%7Zq#^$xYtYYhKUB$~RF|AyxzML&$ zUR5#&Of3acsm)uksyE6Iu;utTM&8DSJetpJ%*_a%B-seFJlcfu5b4T<2dc~u^q+!ha6=tiu??a#gLuDP#q)uH*p6m8z?(Ro5oh?inOef>`mpe*5< zTQNexKe11p142LVX;{02gy4C?7Xk3mqyEGfjSIa%gtq6;cum_2Y%lnXslR>7uwI1G zu;P_l2ehK)HpB|aRvsbpDWA>|Ej#AhQJ=7Nx^5Lz4Pdu#F&M~?;#fVou_*hOpcd7) zJFW(0VMV4{$4pSfkIH?$s6whhO2VIzGrk9c8oYkrov;aqLvZYzvd`e2n&ZInKBB{d z3zTe!?R2;z4_YU8JcZNSF}{ty`WCfaM83n+H%Z3UEJgEq^LAer@_4B3{03cHLY{N- z){8N|P(E|{g~_D#65zmD2cv5?1&(DFbzEzkTE#TO;w8hTm%4)C-_A`zHGk*7n6jwk z|IrryRkmzfTgI2!^1G!imuqmP;Y=DB09^L2y0< z03lJ%HbyySL;4{aGl0I_jav|~TGr=@Q}fN2bxo7Uu%TgUX~J?hzyCTt%ad)U=WCK; z|1vEfEwyx-iU3TJ%hWfBTn91kNje17m2{XXN)h!Y5rW80Xp7KzIZii4g>J(6Rjcbi zNk&;w>dry<9!B+FLHf?bJ08@ydOmna7hA!u;yNn$RcDL7xQ2py)Z>o5)CM&> zbaYfV=&SBNj%uzR5|F-$i}ImoIl;889)jyV!wP2%2o`=S6)1IX_=g-6sPoxgK#eNFe4bBlQP0Oux1hC^8h>wox|Txh}^w+Nk{NYl)d;~2+Y zT6Vpv9PXHKs7hRol1$Mx->@UY1~0_gL^uuOgEYgn_W^{aCmi#5#J}yXGwyn9i*dfj zJ=Vfm7jX4UNFhW&z~?;_{B3-9b~aBFF) zY2mk0eA@o^qGkP)UM_wN%BOhr7ZAb{c4Y0_C|Sf}XWy|c;RyGYweJc~)KGe&E*dCn zVnH-f*2N96h_WG;#4^eSaZ{|IY>IcpEtEHey=VEi&oR5TP3Mn4CShi4Cuz6a5_a_hyd z->`{2VHq}%vorR4_R`(6{$k-kP8=q_Q?QMX3Ogw5j-m;S`1)mQkjAp5Lm>uX>o`@1 zdKh$NYZptWNf5PuCAADxcd#1_QH9q)wDz)bB!9RYV28TZO_L+3GAUY_LQT7y3d|=X z+W3)dm!n&;4D?9Jjp6v~=4OO5cGGzHxIqUf4SM>Z4WpN~!5pwn8@AclG+5}zP%Z`) z&9a#+<+H-kqquNR!y$C0LOchv^C?=PSiNp)_hN0!RX3skPXK>*4ZxW^$#CTVJH+O+ zzOCp$zN5%76gjA>lc=CJrcev0qtjkk{o$NxrW<~%-4`!9l$KXFXnq8%R}x8rEg~F<5c+l>nDXv-AZ_!nT3m57*`FYxs8|Kp0hNy6OOkqQEBG7T7_Pj znM_=<*ekq%Hel+t3sX^3Y=mpTp%!R!BfP2}8jV=BCXT4(4pyqK z0YormSzbWe82uevsqFYP1XyznIVBOz2k9|CmV{);gG9j0yccG~m?@V5ATp(3WEO$;+r-;SR z$<5D6tI$@274cadoRqQb&M;TE(YRmhbp9Lgv}*$Xo9M3zWUDb?igHB7Z$t~^FebwO z2t`b@#sSimV24bt=k{xs+Xp6*nG(PX_l+~L%^Tuy;$+Tiw{N|%C+^9!%*&lK`-$~p zbK(J1Pk6s$rOdq^IgW)9HH@g`Zr?^(&uck*!^o-n_Pb`5m%sSj%3ZWIa`%n>o;7iE z1)eA-U#bcDn$1?KDqBaL->T!g)zR3MXT&on7^YO?Bd z&hs@GUELu;9-xTjcKn6HiDXiEdVow@)Rin1mJAEe?6>e#ko0A70~xoA8$^PqSg7s` zhK~Ju;l_bJESxaz7cL#GSe)*-jWg)Nr2$2KT32{cFqY~nX|G;Q#fI^*O-&?XHC%hy zzQ>w&lcBWlJIhRci2CY+LOG~nm#^?b>|)%@hI>88Ew~rn4XslNHxLfo%k&}^nvj#Y z)CUg6PM*NQ;GA%9*{{86=pm%5(7SI!&!+I43q9}jPn8@08vgyutzAxWmZZn)NQY6_AHWMK0LJ{M z5w4BFldTUl5947ZW5hr}H`iGh?M8?^5VS_zfuf;(hOp5gYK2tc5xO5q&4ZGrhjiZr zqx0T~Zc)`naM#q;wCD8EU;y1Qp#&TagoK{rBn)=zdkUs72(M3%+A|gLu*Vw?v%t^;7Enb7i&}b$kC3sWsO)dBE9~vU zeqPvnh5ghoujJ-tXznd0gCYE8*Gx)?#Uc*V*7)woWPgk=fiS)rgG)2q&?+&gq_CRH z8VJ4eSwXqN)F){57Yh$OIUE}X*uSX{&CP71Cj)XrLj&e>;mY0&|Cu0GguIz+ZZC!z jh-r$(W}J#qBp(y?Ys&czdyzHj_%|C%jkU&d;~W10E;gP~ literal 0 HcmV?d00001 diff --git a/train_code/selective_search/__pycache__/structure.cpython-37.pyc b/train_code/selective_search/__pycache__/structure.cpython-37.pyc new file mode 100644 index 0000000000000000000000000000000000000000..554e29803f5cbf39ff637ffa3608c554a7b6da83 GIT binary patch literal 4319 zcmbtY&2Qtz6`vVW6fMgi@p`@M%}1Thwu#cz4q9Lv6pp*yrcF+Ukqxkg(14&dW0^KZ z$}`liEl@>K2j~Z|hu(9rkNp$+pAAL+pXzkbkCJYO z6h#lU4$99_%m=7)p0j}E4r-o@3vtOnIhlJ=%h{20$uB%RQkM+e`EJWEoT%3=nkMQ6 zCevZ2x@U#kN&9C-9TS03-9jY&!W$&f$v_oemd26E3NH@#Rcu;r;TjbmXdmy^^Hj~D z`U}&qq50Nq9i~0CWg;mLqpjzuJ~sWZqqg>xHfa*ZTc0bV=x}>sA4BMda_dQU7ONlb zg;A23txlReQ#w;}3&KT7&`BlusTj8tJyQ^mc5gxQK_?q%wF$ZXAc&GE3j&hFkYJ2^ z_%-7Fh<}9*J{4;-`NiZGOVO{Ow49awR_&7qv?^ zax#zlxj4o;Dc0E^Syul&JK)wTd|CU1opamwU?t=J9gL0K$-Yr7^K-f_aV|z)HV2x6 zs^d=G+{^2-{xz*-x!7hWcb~!~aFhP7u0c*sC#a4=W=aNf3&pIV>f#t*E{^ygU?%He z@n>=3J}IvXMqV>nsI!(R#O8+ukH{Ca>PBf2m^kVvd!u>qoWpA%^RRUkkR1oj?gQypFRcZH2^b#`u z!ZCxMCY=|xY|xKYdljec3}V`J{fhS5_2|^ge2s{R7*X=?7B%Z-`X2Frow&Na&>ZVG zsQ)HabiBz?>1v9FDRq-zhE|gYG}WHi(~V)fr=YV=P8x~8^se9br1JbZHk3!Z6b{6x3H+Z7ppgl0xj1n)>=j}cfFGz zl?GB-NwIQMe?dQtw0QsqhLlT6gxc`)ZSnZ<-B+rv^U5E*HL=LjGK zATPv7Wa5Ifgz&x25U^beE4c$FcZPp!Y}|PELLK&0hQL`3|DbxRA=+TgXyYdv2&q%! z*K_H?Ms}z+N>i`ae6HXK$-l9$=4pSWP4Fk5o|?2Tz~Bb>mxw`BwWV5H)Y1VCf|%p> z5aC-jO$sXkrAgERRg0)5SNS&R!XHEVnSPh%?KM_+slSV{4^c`CXwtQ2W4GJ56JX7$l$`Gi!hstCDtUrp8*korn<7~+Xs&$Ul+0{SsSeR%zg=bQN zR!tUO)S^uFkP(%hGOfyDZT4#(W223yj1Ry6KZw^MWTl6kvkE60PEqgw6~%H}nyvBS$;j3BJgm0)yCor`LmIF#d{FwVy7aed=) zP&dZCY8=$LIBOfo>zJ``U!|SYvr?==j@)~l=KQG@)SFr%=e~AFe(u|oIomdC zJ0j!<;vXXfQWiJvBcIqVk-jmuoeC4|WnPh?U6EF=!2zB^m{Fof1sN@`aC+fs8A`pd z-&cvWArr$8=_l0KWh?a~~CR7J1WI6U+Rr{k6!Yno~m8w22rl?_+lI(Nk-5 z7uY#5K+uBdN%~wXeaD4o8U16tc&t~gdg(I?V}$1xcB~@X3s@LrX&{loBHnG$>`Sh8 ztK56-%=k{BBN<}F9G5TSw@3)?H2LtonM#_$q6$l-CrWND0_}mAxZD(SKeWibiG=o5 zi0#8nxkc=I*JwSbAA+-0h7H5|d(@qQ%vOJsm~I21jpSr?FDvGqDVQGzy5In58yVm@ zNy9!vFu@3?1%3$n!0C@>s_DOASgDPGfSk;#jk1A>ahXN>$drNrdiRT8A_;wfS^ZN~ z{~?(AkqK0uDW~%DO;@#E)ePboaJ%XTU1OX6x5EI=Df|EfQmou6^JM=qLwWuOA)<2?ZhpL}Z z^_VI;il)GDSIVYfmLvnBg;~Oesit#e(TX{j<$PmPtIhL{aal&>-zEyz=AU@|`6Tu(>?(4%vjE Ne=~gKh4k~Ce*qDm=ji|d literal 0 HcmV?d00001 diff --git a/train_code/selective_search/__pycache__/util.cpython-37.pyc b/train_code/selective_search/__pycache__/util.cpython-37.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f295ac137ef21a7502f58018b9e7487190c21366 GIT binary patch literal 2012 zcmZWpPj6f`5PyDtZ};tH)23;JND->xLRX5ip%M~81*%qRQ&FObAP#$poM-P_v+vJh zdlNS7UMO(u6$#V>+<*_jov-05r+fuY{QS1N3B{7fGvgW0c>MD`SYK}k7#r^XahU|c z-(Fab5X^l{a|j{=LGqwau!eb9&^{$zNAswN`*F=jd826dn?=%3idMe`T`U^k1^u>Y ziUjMLXo)t~jvz-tckKeF1_xc@nqyswNov7%1~NaFr8z(2`AL_$rp`|84$bL;$@z&( zye89fn)>wopqLHJhnQvy#0IlqMg$R|pwFXuI13Rj;!jBswAst?Se8a;>U~ zVa9S@E40j1RT?%ie%lTuE2JG(f>nc+|D!8;?yj^WCIz}YQzcI3cIC=h&UGd=8?ls6jP#ykEeBG|GL>vl>7_wh72~Rur8Sp`wK|12!!}rWvb*k8}qT$m!5?xkBvH)aDI1hkJS;Gt=K2BD&H~e zQ*KdpriyVc3&gCmUwCIf|K!KT#}@)c$799yO|kf7T(Ug5OSFF?7*JJYj*MIr>#>p9 z?}5ICoqx*SuqtHFC?STbcUI|<8S_;34y85~a_)U5jpRBV9{Su3L@s(qcACq(hg_i$ zy|gM%rM6P^tcK<|tpv`K`Qj}Q%X{}GR^_|nsoSVvjQFby0r`Q6K!Vmrs5h{#(B)AU zW7g>U6O2dJ2u2=d&j4h|A!KIBN~VzaY`Bsm$Op5)(l4V~{VAS@Zv=Dr*WOL+{_y|Z zg+J^Qk$fM%q^$wHxAben@3%!3|FMY&bQ>p&wO`^H`IUT2L<0PG&i`6r?X$-Z{1XH3 znY1jea)7I@ICx)21B6~8USwCUGN+*6U+x6=lE}{19#Z|(H9_#IQtX1Jc8}BfvyU_UFYs7I7boDktdF#6KMSD z0DU9kUn}D^PYy2pi#f7?e)F3-Q-+n5Wo9z2JUVu}4eg&<`xn)|X}GW)J8C9{Yv79& z+;&5~=Jkz5b!vbv$aSyYG-t}DLjVj#JFbiN8uW&N1kHpb;pN;U5u{CWt6pzW{U&Tu zpRpf4g%)NFtUtgAk^xGX6f=53YdO9mH=qg8)M8L7nxNhaT0DO%i~Npgi#5^7V!toe z#WiTX&Gi@2oOq~M@+{P+b$i^n7l;jj5U|}2#EtDGx4GhOZlkw`!91MPc?39yGdhc& zQGjXdKsS-Gz5(K*ho2uk!tV}8#x?4?7#A5;7`=(@i@P9w3X_W;!KoeGU*80$eckly zAdWn8tV}s6NWj1{e-sWwPAdW~x61t!b(jaZxCLOXBwMl{} zy%XwPA8Fmi<)j!-^)@8eQWlMSNl!f5Z(H6t1j9rz^uDajhH&h gm7IOc)HN8Bd${5xmmhgdnGQ%oIw3No$+hJ6zkAXK$^ZZW literal 0 HcmV?d00001 diff --git a/train_code/selective_search/core.py b/train_code/selective_search/core.py index 502f9e9..129e8db 100755 --- a/train_code/selective_search/core.py +++ b/train_code/selective_search/core.py @@ -1,7 +1,7 @@ from joblib import Parallel, delayed from skimage.segmentation import felzenszwalb -from util import oversegmentation, switch_color_space, load_strategy -from structure import HierarchicalGrouping +from .util import oversegmentation, switch_color_space, load_strategy +from .structure import HierarchicalGrouping diff --git a/train_code/selective_search/structure.py b/train_code/selective_search/structure.py index 636bf0d..fe64462 100755 --- a/train_code/selective_search/structure.py +++ b/train_code/selective_search/structure.py @@ -2,7 +2,7 @@ from skimage.segmentation import find_boundaries from skimage.segmentation import felzenszwalb from scipy.ndimage import find_objects -import measure +from .measure import * class HierarchicalGrouping(object): @@ -14,7 +14,7 @@ def __init__(self, img, img_seg, sim_strategy): def build_regions(self): self.regions = {} - lbp_img = measure.generate_lbp_image(self.img) + lbp_img = generate_lbp_image(self.img) for label in self.labels: size = (self.img_seg == 1).sum() region_slice = find_objects(self.img_seg==label)[0] @@ -22,8 +22,8 @@ def build_regions(self): [region_slice[i].stop for i in (1,0)]) mask = self.img_seg == label - color_hist = measure.calculate_color_hist(mask, self.img) - texture_hist = measure.calculate_texture_hist(mask, lbp_img) + color_hist = calculate_color_hist(mask, self.img) + texture_hist = calculate_texture_hist(mask, lbp_img) self.regions[label] = { 'size': size, @@ -39,7 +39,7 @@ def build_region_pairs(self): neighbors = self._find_neighbors(i) for j in neighbors: if i < j: - self.s[(i,j)] = measure.calculate_sim(self.regions[i], + self.s[(i,j)] = calculate_sim(self.regions[i], self.regions[j], self.img.size, self.sim_strategy) @@ -118,7 +118,7 @@ def calculate_similarity_for_new_region(self): for j in neighbors: # i is larger than j, so use (j,i) instead - self.s[(j,i)] = measure.calculate_sim(self.regions[i], + self.s[(j,i)] = calculate_sim(self.regions[i], self.regions[j], self.img.size, self.sim_strategy) diff --git a/train_code/train.py b/train_code/train.py old mode 100644 new mode 100755 index fb3c5bb..1c3571f --- a/train_code/train.py +++ b/train_code/train.py @@ -6,21 +6,20 @@ import tensorflow as tf -import tensorflow.contrib.slim as slim - +import tf_slim as slim import utils import os import numpy as np import argparse import network import loss - +import random from tqdm import tqdm from guided_filter import guided_filter - +random.seed(0) os.environ["CUDA_VISIBLE_DEVICES"]="0" - +tf.compat.v1.disable_eager_execution() def arg_parser(): parser = argparse.ArgumentParser() parser.add_argument("--patch_size", default = 256, type = int) @@ -40,11 +39,11 @@ def arg_parser(): def train(args): - input_photo = tf.placeholder(tf.float32, [args.batch_size, + input_photo = tf.compat.v1.placeholder(tf.float32, [args.batch_size, args.patch_size, args.patch_size, 3]) - input_superpixel = tf.placeholder(tf.float32, [args.batch_size, + input_superpixel = tf.compat.v1.placeholder(tf.float32, [args.batch_size, args.patch_size, args.patch_size, 3]) - input_cartoon = tf.placeholder(tf.float32, [args.batch_size, + input_cartoon = tf.compat.v1.placeholder(tf.float32, [args.batch_size, args.patch_size, args.patch_size, 3]) output = network.unet_generator(input_photo) @@ -62,14 +61,14 @@ def train(args): scale=1, patch=True, name='disc_blur') - vgg_model = loss.Vgg19('vgg19_no_fc.npy') + vgg_model = loss.Vgg19('train_code/vgg19_no_fc.npy') vgg_photo = vgg_model.build_conv4_4(input_photo) vgg_output = vgg_model.build_conv4_4(output) vgg_superpixel = vgg_model.build_conv4_4(input_superpixel) h, w, c = vgg_photo.get_shape().as_list()[1:] - photo_loss = tf.reduce_mean(tf.losses.absolute_difference(vgg_photo, vgg_output))/(h*w*c) - superpixel_loss = tf.reduce_mean(tf.losses.absolute_difference\ + photo_loss = tf.reduce_mean(input_tensor=tf.compat.v1.losses.absolute_difference(vgg_photo, vgg_output))/(h*w*c) + superpixel_loss = tf.reduce_mean(input_tensor=tf.compat.v1.losses.absolute_difference\ (vgg_superpixel, vgg_output))/(h*w*c) recon_loss = photo_loss + superpixel_loss tv_loss = loss.total_variation_loss(output) @@ -77,29 +76,29 @@ def train(args): g_loss_total = 1e4*tv_loss + 1e-1*g_loss_blur + g_loss_gray + 2e2*recon_loss d_loss_total = d_loss_blur + d_loss_gray - all_vars = tf.trainable_variables() + all_vars = tf.compat.v1.trainable_variables() gene_vars = [var for var in all_vars if 'gene' in var.name] disc_vars = [var for var in all_vars if 'disc' in var.name] - tf.summary.scalar('tv_loss', tv_loss) - tf.summary.scalar('photo_loss', photo_loss) - tf.summary.scalar('superpixel_loss', superpixel_loss) - tf.summary.scalar('recon_loss', recon_loss) - tf.summary.scalar('d_loss_gray', d_loss_gray) - tf.summary.scalar('g_loss_gray', g_loss_gray) - tf.summary.scalar('d_loss_blur', d_loss_blur) - tf.summary.scalar('g_loss_blur', g_loss_blur) - tf.summary.scalar('d_loss_total', d_loss_total) - tf.summary.scalar('g_loss_total', g_loss_total) + tf.compat.v1.summary.scalar('tv_loss', tv_loss) + tf.compat.v1.summary.scalar('photo_loss', photo_loss) + tf.compat.v1.summary.scalar('superpixel_loss', superpixel_loss) + tf.compat.v1.summary.scalar('recon_loss', recon_loss) + tf.compat.v1.summary.scalar('d_loss_gray', d_loss_gray) + tf.compat.v1.summary.scalar('g_loss_gray', g_loss_gray) + tf.compat.v1.summary.scalar('d_loss_blur', d_loss_blur) + tf.compat.v1.summary.scalar('g_loss_blur', g_loss_blur) + tf.compat.v1.summary.scalar('d_loss_total', d_loss_total) + tf.compat.v1.summary.scalar('g_loss_total', g_loss_total) - update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS) + update_ops = tf.compat.v1.get_collection(tf.compat.v1.GraphKeys.UPDATE_OPS) with tf.control_dependencies(update_ops): - g_optim = tf.train.AdamOptimizer(args.adv_train_lr, beta1=0.5, beta2=0.99)\ + g_optim = tf.compat.v1.train.AdamOptimizer(args.adv_train_lr, beta1=0.5, beta2=0.99)\ .minimize(g_loss_total, var_list=gene_vars) - d_optim = tf.train.AdamOptimizer(args.adv_train_lr, beta1=0.5, beta2=0.99)\ + d_optim = tf.compat.v1.train.AdamOptimizer(args.adv_train_lr, beta1=0.5, beta2=0.99)\ .minimize(d_loss_total, var_list=disc_vars) ''' @@ -107,28 +106,32 @@ def train(args): config.gpu_options.allow_growth = True sess = tf.Session(config=config) ''' - gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=args.gpu_fraction) - sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options)) + gpu_options = tf.compat.v1.GPUOptions(per_process_gpu_memory_fraction=args.gpu_fraction) + sess = tf.compat.v1.Session(config=tf.compat.v1.ConfigProto(gpu_options=gpu_options)) - train_writer = tf.summary.FileWriter(args.save_dir+'/train_log') - summary_op = tf.summary.merge_all() - saver = tf.train.Saver(var_list=gene_vars, max_to_keep=20) + train_writer = tf.compat.v1.summary.FileWriter(args.save_dir+'/train_log') + summary_op = tf.compat.v1.summary.merge_all() + saver = tf.compat.v1.train.Saver(var_list=gene_vars, max_to_keep=20) with tf.device('/device:GPU:0'): - sess.run(tf.global_variables_initializer()) - saver.restore(sess, tf.train.latest_checkpoint('pretrain/saved_models')) + sess.run(tf.compat.v1.global_variables_initializer()) + saver.restore(sess, tf.train.latest_checkpoint('pretrainsave_models')) - face_photo_dir = 'dataset/photo_face' + face_photo_dir = 'dataset/face_photo' face_photo_list = utils.load_image_list(face_photo_dir) - scenery_photo_dir = 'dataset/photo_scenery' + scenery_photo_dir = 'dataset/scenery_photo' scenery_photo_list = utils.load_image_list(scenery_photo_dir) - face_cartoon_dir = 'dataset/cartoon_face' - face_cartoon_list = utils.load_image_list(face_cartoon_dir) - scenery_cartoon_dir = 'dataset/cartoon_scenery' - scenery_cartoon_list = utils.load_image_list(scenery_cartoon_dir) + face_cartoon_dir_kyoto_face = 'dataset/face_cartoon/kyoto_face/' + face_cartoon_list = utils.load_image_list(face_cartoon_dir_kyoto_face) + face_cartoon_dir_pa_face = 'dataset/face_cartoon/pa_face/' + face_cartoon_list.extend(utils.load_image_list(face_cartoon_dir_pa_face)) + scenery_cartoon_dir = 'dataset/scenery_cartoon/' + scenery_cartoon_list = utils.load_image_list(scenery_cartoon_dir+"hayao/") + scenery_cartoon_list.extend(utils.load_image_list(scenery_cartoon_dir+"hosoda/")) + scenery_cartoon_list = utils.load_image_list(scenery_cartoon_dir+"shinkai/") for total_iter in tqdm(range(args.total_iter)): @@ -172,10 +175,12 @@ def train(args): print('Iter: {}, d_loss: {}, g_loss: {}, recon_loss: {}'.\ format(total_iter, d_loss, g_loss, r_loss)) - if np.mod(total_iter+1, 500 ) == 0: + if np.mod(total_iter+1, 250 ) == 0: saver.save(sess, args.save_dir+'/saved_models/model', write_meta_graph=False, global_step=total_iter) - + + if np.mod(total_iter+1, 500 ) == 0: + photo_face = utils.next_batch(face_photo_list, args.batch_size) cartoon_face = utils.next_batch(face_cartoon_list, args.batch_size) photo_scenery = utils.next_batch(scenery_photo_list, args.batch_size) @@ -193,15 +198,14 @@ def train(args): str(total_iter)+'_face_result.jpg', 4) utils.write_batch_image(photo_face, args.save_dir+'/images', str(total_iter)+'_face_photo.jpg', 4) - utils.write_batch_image(result_scenery, args.save_dir+'/images', str(total_iter)+'_scenery_result.jpg', 4) utils.write_batch_image(photo_scenery, args.save_dir+'/images', str(total_iter)+'_scenery_photo.jpg', 4) + if __name__ == '__main__': args = arg_parser() train(args) - \ No newline at end of file diff --git a/train_code/utils.py b/train_code/utils.py old mode 100644 new mode 100755 index 472559c..96f9cf3 --- a/train_code/utils.py +++ b/train_code/utils.py @@ -167,4 +167,3 @@ def write_batch_image(image, save_dir, name, n): if __name__ == '__main__': pass - \ No newline at end of file From 25452bc7e9803e5c8f3d967773d051f407036a45 Mon Sep 17 00:00:00 2001 From: Siddhi Vinayak Tripathi Date: Mon, 2 Nov 2020 17:34:24 +0530 Subject: [PATCH 2/3] Disabled eager execution --- test_code/cartoonize.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test_code/cartoonize.py b/test_code/cartoonize.py index 4a433c4..533b0b1 100755 --- a/test_code/cartoonize.py +++ b/test_code/cartoonize.py @@ -6,7 +6,7 @@ import guided_filter from tqdm import tqdm - +tf.compat.v1.disable_eager_execution() def resize_crop(image): h, w, c = np.shape(image) From fbc572042bec2c3663a4edae0e7998202ef99f5a Mon Sep 17 00:00:00 2001 From: Siddhi Vinayak Tripathi Date: Wed, 4 Nov 2020 17:36:19 +0530 Subject: [PATCH 3/3] Removed use of wandb api --- train_code/pretrain.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/train_code/pretrain.py b/train_code/pretrain.py index 0f694af..29c9f58 100755 --- a/train_code/pretrain.py +++ b/train_code/pretrain.py @@ -82,7 +82,7 @@ def train(args): if np.mod(total_iter+1, 50) == 0: - wandb.log({"r_loss":r_loss,"iteration":total_iter}) + print('pretrain, iter: {}, recon_loss: {}'.format(total_iter, r_loss)) if np.mod(total_iter+1, 500 ) == 0: