twitter_compression/encoder.py at master · yurun00/twitter_compression · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
import tensorflow as tf
import numpy as np
from common import *
from ConvResCell import ConvResCell

def build_encoder(input_images, hps=None):
  with tf.variable_scope('encoder_net') as scopes:
    batch_size = input_images.get_shape().as_list()[0]

    with tf.variable_scope('enc_in_layer') as scope:
      conv = add_conv2d(input_images, 5, 3, 64, [1,2,2,1], 'enc_in_conv')
      hidden = add_relu(conv, hps['relu_leakiness'], 'enc_in_hidden')

    res_1 = ConvResCell(kernel_sizes=[5, 3, 3], strides=[2, 1, 1], channels=[128, 128, 128],
                        scope='res_1', hps=hps)
    output_flow = res_1(hidden)

    res_2 = ConvResCell(kernel_sizes=[3, 3], strides=[1, 1], channels=[128, 128],
                        scope='res_2', hps=hps)
    output_flow = res_2(output_flow)

    with tf.variable_scope('enc_out_layer') as scope:
      conv = add_conv2d(output_flow, 5, output_flow.get_shape().as_list()[3], hps['codec_dim'], [1,2,2,1], 'enc_out_conv')
      hidden = add_relu(conv, hps['relu_leakiness'], 'enc_out_hidden')
      # hidden = add_tanh(conv, 'enc_out_hidden')

    tf.summary.histogram('codec', hidden)

    return hidden


def build_model():
  with tf.Graph().as_default():
    input_image = tf.Variable(tf.zeros([1, 128, 128, 3]), dtype=tf.float32, name='input_image')

    code = build_encoder(input_image, hps=HParams)

    sess = tf.Session()
    init = tf.global_variables_initializer()
    sess.run(init)

    summary_writer = tf.summary.FileWriter('log/build_model/encoder/', sess.graph)

    merged = tf.summary.merge_all()
    summary = sess.run(merged)
    summary_writer.add_summary(summary, 1)

def main(_):
  build_model()

if __name__ == '__main__':
  tf.app.run()