Class Model

examples/SuperResolution/enet-pat.py:45–218

class Model(GANModelDesc):

    def __init__(self, height=SHAPE_LR, width=SHAPE_LR):
        super(Model, self).__init__()
        self.height = height
        self.width = width

    def inputs(self):
        return [tf.TensorSpec((None, self.height * 1, self.width * 1, CHANNELS), tf.float32, 'Ilr'),
                tf.TensorSpec((None, self.height * 4, self.width * 4, CHANNELS), tf.float32, 'Ihr')]

    def build_graph(self, Ilr, Ihr):
        Ilr, Ihr = Ilr / 255.0, Ihr / 255.0
        Ibicubic = tf.image.resize_bicubic(
            Ilr, [4 * self.height, 4 * self.width], align_corners=True,
            name='bicubic_baseline')    # (0,1)

        VGG_MEAN_TENSOR = tf.constant(VGG_MEAN, dtype=tf.float32)

        def resnet_block(x, name):
            with tf.variable_scope(name):
                y = Conv2D('conv0', x, NF, activation=tf.nn.relu)
                y = Conv2D('conv1', y, NF, activation=tf.identity)
            return x + y

        def upsample(x, factor=2):
            _, h, w, _ = x.get_shape().as_list()
            x = tf.image.resize_nearest_neighbor(x, [factor * h, factor * w], align_corners=True)
            return x

        def generator(x, Ibicubic):
            x = x - VGG_MEAN_TENSOR / 255.0
            with argscope(Conv2D, kernel_size=3, activation=tf.nn.relu):
                x = Conv2D('conv1', x, NF)
                for i in range(10):
                    x = resnet_block(x, 'block_%i' % i)
                x = upsample(x)
                x = Conv2D('conv_post_1', x, NF)
                x = upsample(x)
                x = Conv2D('conv_post_2', x, NF)
                x = Conv2D('conv_post_3', x, NF)
                Ires = Conv2D('conv_post_4', x, 3, activation=tf.identity)
                Iest = tf.add(Ibicubic, Ires, name='Iest')
                return Iest     # [0,1]

        @auto_reuse_variable_scope
        def discriminator(x):
            x = x - VGG_MEAN_TENSOR / 255.0
            with argscope(Conv2D, kernel_size=3, activation=tf.nn.leaky_relu):
                x = Conv2D('conv0', x, 32)
                x = Conv2D('conv0b', x, 32, strides=2)
                x = Conv2D('conv1', x, 64)
                x = Conv2D('conv1b', x, 64, strides=2)
                x = Conv2D('conv2', x, 128)
                x = Conv2D('conv2b', x, 128, strides=2)
                x = Conv2D('conv3', x, 256)
                x = Conv2D('conv3b', x, 256, strides=2)
                x = Conv2D('conv4', x, 512)