def encoder(input, patch_size, quan_scale): # input = tf.Print(input, [tf.shape(input)[0], tf.shape(input)[1], tf.shape(input)[2], tf.shape(input)[3]], 'Input tensor shape:') # input = tf.Print(input, [input], 'Input tensor:') with tf.variable_scope('normalize'): output = tf.reshape(input, [-1, patch_size, patch_size, 3]) tf.summary.image("ori_image", output, max_outputs=4) tf.summary.histogram('ori_input', output) output = (output - train_data_mean) / train_data_std tf.summary.histogram('normalized_input', output) # output = tf.Print(output, [output], 'Normalized tensor:') output = basic_block.my_conv2d( inputs=output, filters=32, strides=[2, 2], kernel_size=[3, 3], padding="SAME", activation=tf.nn.relu, kernel_initializer=tf.random_normal_initializer(mean=0.0, stddev=0.01), bias_initializer=tf.constant_initializer(0.0), name='encode_0') output = basic_block.my_conv2d( inputs=output, filters=32, strides=[2, 2], kernel_size=[3, 3], padding="SAME", activation=tf.nn.relu, kernel_initializer=tf.random_normal_initializer(mean=0.0, stddev=0.01), bias_initializer=tf.constant_initializer(0.0), name='encode_1') output = basic_block.my_conv2d( inputs=output, filters=64, strides=[2, 2], kernel_size=[3, 3], padding="SAME", activation=tf.nn.relu, kernel_initializer=tf.random_normal_initializer(mean=0.0, stddev=0.01), bias_initializer=tf.constant_initializer(0.0), name='encode_2') output = basic_block.my_conv2d( inputs=output, filters=64, strides=[2, 2], kernel_size=[3, 3], padding="SAME", activation=tf.nn.relu, kernel_initializer=tf.random_normal_initializer(mean=0.0, stddev=0.01), bias_initializer=tf.constant_initializer(0.0), name='encode_3') output = basic_block.res_block( inputs=output, filters=64, strides=[1, 1], kernel_size=[3, 3], padding="SAME", activation=tf.nn.relu, kernel_initializer=tf.random_normal_initializer(mean=0.0, stddev=0.01), bias_initializer=tf.constant_initializer(0.0), name='encode_res_1', layer_num=2) output = basic_block.res_block( inputs=output, filters=64, strides=[1, 1], kernel_size=[3, 3], padding="SAME", activation=tf.nn.relu, kernel_initializer=tf.random_normal_initializer(mean=0.0, stddev=0.01), bias_initializer=tf.constant_initializer(0.0), name='encode_res_2', layer_num=2) output = basic_block.my_conv2d( inputs=output, filters=64, strides=[1, 1], kernel_size=[3, 3], padding="SAME", activation=tf.identity, kernel_initializer=tf.random_normal_initializer(mean=0.0, stddev=0.01), bias_initializer=tf.constant_initializer(0.0), name='encode_4') with tf.variable_scope('sigmoid_scale_quantize'): output = tf.nn.sigmoid(output) * (quan_scale - 1) output = tf.stop_gradient(tf.round(output) - output) + output tf.add_to_collection('bitrate_reg_var', output) tf.summary.histogram('sigmoid_scale_quantize_result', output) return output
def decoder(input, quan_scale): with tf.variable_scope('reverse_sigmoid_scale'): # output = basic_block.reverse_sigmoid(input / 255.0) # stable reverse sigmoid output = basic_block.reverse_sigmoid( (input + 1e-6) / (quan_scale - 1 + 1e-5)) tf.summary.histogram('reverse_sigmoid_scale_result', output) # output = tf.Print(output, [output[5, 20:30, 20, 0]]) output = basic_block.my_conv2d( inputs=output, filters=64, strides=[1, 1], kernel_size=[3, 3], padding="SAME", activation=tf.identity, kernel_initializer=tf.random_normal_initializer(mean=0.0, stddev=0.01), bias_initializer=tf.constant_initializer(0.0), name='decode_4') output = basic_block.res_block( inputs=output, filters=64, strides=[1, 1], kernel_size=[3, 3], padding="SAME", activation=tf.nn.relu, kernel_initializer=tf.random_normal_initializer(mean=0.0, stddev=0.01), bias_initializer=tf.constant_initializer(0.0), name='decode_res_1', layer_num=2) output = basic_block.res_block( inputs=output, filters=64, strides=[1, 1], kernel_size=[3, 3], padding="SAME", activation=tf.nn.relu, kernel_initializer=tf.random_normal_initializer(mean=0.0, stddev=0.01), bias_initializer=tf.constant_initializer(0.0), name='decode_res_2', layer_num=2) output = basic_block.my_conv2d_transpose( inputs=output, filters=64, strides=[2, 2], kernel_size=[3, 3], padding="SAME", activation=tf.nn.relu, kernel_initializer=tf.random_normal_initializer(mean=0.0, stddev=0.01), bias_initializer=tf.constant_initializer(0.0), name='decode_3') output = basic_block.my_conv2d_transpose( inputs=output, filters=32, strides=[2, 2], kernel_size=[3, 3], padding="SAME", activation=tf.nn.relu, kernel_initializer=tf.random_normal_initializer(mean=0.0, stddev=0.01), bias_initializer=tf.constant_initializer(0.0), name='decode_2') # output = tf.Print(output, [output[5, 20:30, 20, 0]]) output = basic_block.my_conv2d_transpose( inputs=output, filters=32, strides=[2, 2], kernel_size=[3, 3], padding="SAME", activation=tf.nn.relu, kernel_initializer=tf.random_normal_initializer(mean=0.0, stddev=0.01), bias_initializer=tf.constant_initializer(0.0), name='decode_1') output = basic_block.my_conv2d_transpose( inputs=output, filters=3, strides=[2, 2], kernel_size=[3, 3], padding="SAME", activation=tf.identity, kernel_initializer=tf.random_normal_initializer(mean=0.0, stddev=0.01), bias_initializer=tf.constant_initializer(0.0), name='decode_0') # output = tf.Print(output, [output[5, 20:30, 20, 0]]) with tf.variable_scope('denormalize'): output = output * train_data_std + train_data_mean tf.summary.histogram('denormalized_input', output) tf.summary.image("recons_image", output, max_outputs=4) # To be modified # output = tf.clip_by_value(output, 0.0, 1.0) output = tf.clip_by_value(output, 0.0, 255.0) # output = tf.Print(output, [output[5, 20:30, 20, 0]]) return output