import tensorflow as tf from utils import logger import ops class Encoder(object): def __init__(self, name, is_train, norm='instance', activation='leaky', image_size=128, latent_dim=8, use_resnet=True): logger.info('Init Encoder %s', name) self.name = name self._is_train = is_train self._norm = norm self._activation = activation self._reuse = False self._image_size = image_size self._latent_dim = latent_dim self._use_resnet = use_resnet def __call__(self, input): if self._use_resnet: return self._resnet(input) else: return self._convnet(input) def _convnet(self, input): with tf.variable_scope(self.name, reuse=self._reuse): num_filters = [64, 128, 256, 512, 512, 512, 512] if self._image_size == 256: num_filters.append(512) E = input for i, n in enumerate(num_filters): E = ops.conv_block(E, n, 'C{}_{}'.format(n, i), 4, 2, self._is_train, self._reuse, norm=self._norm if i else None, activation='leaky') E = ops.flatten(E) mu = ops.mlp(E, self._latent_dim, 'FC8_mu', self._is_train, self._reuse, norm=None, activation=None) log_sigma = ops.mlp(E, self._latent_dim, 'FC8_sigma', self._is_train, self._reuse, norm=None, activation=None) z = mu + tf.random_normal(shape=tf.shape(self._latent_dim)) * tf.exp(log_sigma) self._reuse = True self.var_list = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, self.name) return z, mu, log_sigma def _resnet(self, input): with tf.variable_scope(self.name, reuse=self._reuse): num_filters = [128, 256, 512, 512] if self._image_size == 256: num_filters.append(512) E = input E = ops.conv_block(E, 64, 'C{}_{}'.format(64, 0), 4, 2, self._is_train, self._reuse, norm=None, activation='leaky', bias=True) for i, n in enumerate(num_filters): E = ops.residual(E, n, 'res{}_{}'.format(n, i + 1), self._is_train, self._reuse, norm=self._norm, bias=True) E = tf.nn.avg_pool(E, [1, 2, 2, 1], [1, 2, 2, 1], 'SAME') E = tf.nn.relu(E) E = tf.nn.avg_pool(E, [1, 8, 8, 1], [1, 8, 8, 1], 'SAME') E = ops.flatten(E) mu = ops.mlp(E, self._latent_dim, 'FC8_mu', self._is_train, self._reuse, norm=None, activation=None) log_sigma = ops.mlp(E, self._latent_dim, 'FC8_sigma', self._is_train, self._reuse, norm=None, activation=None) z = mu + tf.random_normal(shape=tf.shape(self._latent_dim)) * tf.exp(log_sigma) self._reuse = True self.var_list = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, self.name) return z, mu, log_sigma