def aux_classifier(self, inputs, labels, input_channels, is_training, scope=None):
"""
Auxiliary Classifier used in Inception Module to help propagate
gradients backward.
"""
with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
# pooling layer with 5x5 kernel and stride 3 (new size: 4x4xC)
network = tf.nn.avg_pool(inputs, 5, 3, 'VALID', name='pool')
# convolution with 1x1 kernel and stride 1 (new size: 4x4x128)
network = ops.convolution(network, input_channels, 128, 1, 128, batch_norm=False,
is_training=is_training, scope='auxconv')
# flatten (new size: 2048)
network = ops.flatten(network, scope='flatten')
# fully connected layer (new size: 1024)
network = ops.dense(network, 2048, 1024, dropout=True, dropout_rate=0.7,
is_training=is_training, scope='fc1')
# output layer (new size: 10) -- Original Paper Size: 1000 (for ImageNet)
network = ops.dense(network, 1024, 10, activation=None, is_training=is_training,
scope='fc2')
# loss of auxiliary classifier
loss = ops.loss(network, labels, scope='auxloss')
return loss
def _build_posterior(self, encoding):
'''
Build layers for posterior parameter approximation.
Args:
- encoding: a tensor containing the learned representation of
model inputs that will be linearly transformed to
approximate the parameters to our latent posterior
Returns:
- z_sample: a sampled tensor from the approximated distribution
conditioned on the input.
'''
# Reshape the encoding to shape [batch_size, flattened_dim]
h, w, d = encoding.get_shape().as_list()[1:]
flat_dim = h * w * d
z_dim = self._z_dim
flat_encoding = tf.reshape(encoding, shape=[-1, flat_dim])
# Dense layers for approximating the latent mean and log stddev
# We approximate log sttdev instead of stddev to ensure nonnegativity
self.z_mu = dense(z_dim, flat_encoding)
self.z_log_sigma = 0.5 * dense(z_dim, flat_encoding)
self.z_sigma = tf.exp(self.z_log_sigma)
# Sample from posterior using the 'reparameterization trick'
z_sample = self.z_mu + self.z_sigma * tf.random_normal()
return z_sample
def discriminator(self, x_hr, reuse: bool = False):
""" original paper maybe used VGG-like model, but i just custom-ed it"""
with tf.variable_scope("discriminator", reuse=reuse):
ch: int = self.n_feats
x = x_hr
x = self.res_block_down(x, ch * 1, use_bias=True, sn=self.use_sn,
scope="res_block_down_1")
x = self.res_block_down(x, ch * 2, use_bias=True, sn=self.use_sn,
scope="res_block_down_2")
x = self.res_block_down(x, ch * 4, use_bias=True, sn=self.use_sn,
scope="res_block_down_3")
x = self.res_block_down(x, ch * 8, use_bias=True, sn=self.use_sn,
scope="res_block_down_4")
x = self.res_block_down(x, ch * 16, use_bias=True, sn=self.use_sn,
scope="res_block_down_5")
x = self.res_block(x, ch * 16, use_bias=False, sn=self.use_sn, scope="res_block")
x = tf.nn.leaky_relu(x, alpha=.2)
x = tf.layers.flatten(x)
x = dense(x, units=128, sn=self.use_sn, scope="dense_1")
x = tf.nn.leaky_relu(x, alpha=.2)
x = dense(x, units=1, sn=self.use_sn, scope="disc")
return x
def get_agent(x, reuse=False):
"""
Generate the CNN agent
:param x: tensor, Input frames concatenated along axis 3
:param reuse: bool, True -> Reuse weight variables
False -> Create new ones
:return: Tensor, logits for each valid action
"""
if reuse:
tf.get_variable_scope().reuse_variables()
x = tf.divide(x, 255.0, name='Normalize')
conv_1 = tf.nn.relu(
ops.cnn_2d(x,
weight_shape=mc.conv_1,
strides=mc.stride_1,
name='conv_1'))
conv_2 = tf.nn.relu(
ops.cnn_2d(conv_1,
weight_shape=mc.conv_2,
strides=mc.stride_2,
name='conv_2'))
conv_3 = tf.nn.relu(
ops.cnn_2d(conv_2,
weight_shape=mc.conv_3,
strides=mc.stride_3,
name='conv_3'))
conv_3_r = tf.reshape(conv_3, [-1, 7 * 7 * 64], name='reshape')
dense_1 = tf.nn.relu(
ops.dense(conv_3_r, 7 * 7 * 64, mc.dense_1, name='dense_1'))
output = ops.dense(dense_1, mc.dense_1, mc.dense_2, name='dense_2')
return output
def get_agent(x, reuse=False):
if reuse:
tf.get_variable_scope().reuse_variables()
dense_1 = tf.nn.relu(ops.dense(x, 8, mc.dense_1, name='dense_1'))
dense_2 = tf.nn.relu(ops.dense(dense_1, mc.dense_1, mc.dense_2, name='dense_2'))
dense_3 = tf.nn.relu(ops.dense(dense_2, mc.dense_2, mc.dense_3, name='dense_3'))
output = ops.dense(dense_3, mc.dense_3, env.action_space.n, name='output')
return output
def model(self, x, action, reuse=False):
if reuse:
tf.get_variable_scope().reuse_variables()
# TODO: Use a better network for video frame prediction
# Encoder
x = tf.divide(x, 255.0)
conv_1 = tf.nn.relu(
ops.cnn_2d(x,
weight_shape=[6, 6, 4, 64],
strides=[1, 2, 2, 1],
name='conv_1'))
conv_2 = tf.nn.relu(
ops.cnn_2d(conv_1,
weight_shape=[6, 6, 64, 64],
strides=[1, 2, 2, 1],
name='conv_2',
padding="SAME"))
conv_3 = tf.nn.relu(
ops.cnn_2d(conv_2,
weight_shape=[6, 6, 64, 64],
strides=[1, 2, 2, 1],
name='conv_3',
padding="SAME"))
conv_3_flatten = tf.reshape(conv_3, shape=[-1, 6400], name='reshape_1')
dense_1 = tf.nn.relu(
ops.dense(conv_3_flatten, 6400, 1024, name='dense_1'))
dense_2 = ops.dense(dense_1, 1024, 2048, name='dense_2')
action_dense_1 = ops.dense(action, 4, 2048, name='action_dense_1')
dense_2_action = tf.multiply(dense_2,
action_dense_1,
name='dense_2_action')
# Decoder
dense_3 = ops.dense(dense_2_action, 2048, 1024, name='dense_3')
dense_4 = tf.nn.relu(
ops.dense(dense_3, 1024, 11 * 11 * 64, name='dense_4'))
dense_4_reshaped = tf.reshape(dense_4,
shape=[self.batch_size, 11, 11, 64],
name='dense_4_reshaped')
conv_t_1 = tf.nn.relu(
ops.cnn_2d_trans(dense_4_reshaped,
weight_shape=[6, 6, 64, 64],
strides=[1, 2, 2, 1],
output_shape=[self.batch_size, 21, 21, 64],
name='conv_t_1'))
conv_t_2 = tf.nn.relu(
ops.cnn_2d_trans(conv_t_1,
weight_shape=[6, 6, 64, 64],
strides=[1, 2, 2, 1],
output_shape=[self.batch_size, 42, 42, 64],
name='conv_t_2'))
output = ops.cnn_2d_trans(conv_t_2,
weight_shape=[6, 6, 1, 64],
strides=[1, 2, 2, 1],
output_shape=[self.batch_size, 84, 84, 1],
name='output_image')
return output
def fc_gen(z, units):
"""
Fully-connected Generator.
"""
dense1 = tf.nn.relu(dense(z, units, "gen/dense1"))
dense2 = tf.nn.relu(dense(dense1, units, "gen/dense2"))
dense3 = tf.nn.sigmoid(dense(dense2, 28 * 28, "gen/dense3"))
deflat = tf.reshape(dense3, shape=[-1, 28, 28, 1])
return deflat
def fc_discr(images, units):
"""
Fully-connected Discriminator.
"""
flat = tf.reshape(images, shape=[-1, 28 * 28])
dense1 = tf.nn.relu(dense(flat, units, "discr/dense1"))
dense2 = tf.nn.relu(dense(dense1, units, "discr/dense2"))
dense3 = tf.nn.sigmoid(dense(dense2, 1, "discr/dense3"))
return dense3
def encode(self, input_images):
with tf.variable_scope("encode"):
# 28x28x1 -> 14x14x16
h1 = lrelu(conv2d(input_images, 1, self.e_h1, "e_h1"))
# 14x14x16 -> 7x7x32
h2 = lrelu(conv2d(h1, self.e_h1, self.e_h2, "e_h2"))
h2_flat = tf.reshape(h2, [-1, 7*7*self.e_h2])
w_mean = dense(h2_flat, 7*7*self.e_h2, self.n_z, "w_mean")
w_stddev = dense(h2_flat, 7*7*self.e_h2, self.n_z, "w_stddev")
return w_mean, w_stddev
def build_model(self, inputs, labels, is_training):
# pad inputs to size 224x224x3 - NOTE: may change to bilinear upsampling
pad = int((self.image_size - self.height) / 2)
inputs = tf.pad(inputs, [[0, 0], [pad, pad], [pad, pad], [0, 0]])
# convolution with 11x11 kernel and stride 4 (new size: 55x55x96)
self.network = ops.convolution(inputs, self.channels, 96, 11, 96, stride=4,
padding='VALID', is_training=is_training, scope='conv1')
# pooling with 3x3 kernel and stride 2 (new size: 27x27x96)
self.network = ops.pooling(self.network, k_size=3, scope='pool1')
# convolution with 5x5 kernel and stride 1 (new size: 27x27x256)
self.network = ops.convolution(self.network, 96, 256, 5, 256,
is_training=is_training, scope='conv2')
# pooling with 3x3 kernel and stride 2 (new size: 13x13x256)
self.network = ops.pooling(self.network, k_size=3, scope='pool2')
# convolution with 3x3 kernel and stride 1 (new size: 13x13x384)
self.network = ops.convolution(self.network, 256, 384, 3, 384, batch_norm=False,
is_training=is_training, scope='conv3')
# convolution with 3x3 kernel and stride 1 (new size: 13x13x384)
self.network = ops.convolution(self.network, 384, 384, 3, 384, batch_norm=False,
is_training=is_training, scope='conv4')
# convolution with 3x3 kernel and stride 1 (new size: 13x13x256)
self.network = ops.convolution(self.network, 384, 256, 3, 256, batch_norm=False,
is_training=is_training, scope='conv5')
# pooling with 3x3 kernel and stride 2 (new size: 6x6x256)
self.network = ops.pooling(self.network, k_size=3, scope='pool3')
# flatten (new size: 9216)
self.network = ops.flatten(self.network, scope='flatten')
# fully connected layer (new size: 4096)
self.network = ops.dense(self.network, 9216, 4096, dropout=True, dropout_rate=0.2,
is_training=is_training, scope='fc1')
# fully connected layer (new size: 1024) -- Original Paper Size: 4096 (for ImageNet)
self.network = ops.dense(self.network, 4096, 1024, dropout=True, dropout_rate=0.2,
is_training=is_training, scope='fc2')
# output layer (new size: 10) -- Original Paper Size: 1000 (for ImageNet)
self.network = ops.dense(self.network, 1024, 10, activation=None,
is_training=is_training, scope='fc3')
self.loss = ops.loss(self.network, labels, scope='loss')
if is_training:
self.optimizer = ops.optimize(self.loss, self.learning_rate, scope='update')
def forward(self, state, reuse=False):
with tf.compat.v1.variable_scope('policy'):
h0 = ops.dense(state, self.arch_params['in_dim'],
self.arch_params['n_hidden_0'], tf.nn.relu,
'dense0', reuse)
h1 = ops.dense(h0, self.arch_params['n_hidden_0'],
self.arch_params['n_hidden_1'], tf.nn.relu,
'dense1', reuse)
relu1_do = tf.nn.dropout(h1, self.arch_params['do_keep_prob'])
a = ops.dense(relu1_do, self.arch_params['n_hidden_1'],
self.arch_params['out_dim'], None, 'dense2', reuse)
return a
def tf_buildmodel(x, reuse=None):
x = tf.nn.relu(
conv2D(x, shape=[8, 8, 32], name='1', stride=[1, 4, 4, 1],
reuse=reuse))
x = tf.nn.relu(
conv2D(x, shape=[64, 4, 4], name='2', stride=[1, 2, 2, 1],
reuse=reuse))
x = tf.nn.relu(
conv2D(x, shape=[64, 3, 3], name='3', stride=[1, 1, 1, 1],
reuse=reuse))
x = tf.reshape(x, [-1, 300])
x = dense(x, 512, name='1', reuse=reuse)
x = dense(x, 2, name='prj', activation=None, reuse=reuse)
return x
def forward(self, state, action, reuse=False):
with tf.variable_scope('discriminator'):
concat = tf.concat(axis=1, values=[state, action])
h0 = ops.dense(concat, self.arch_params['in_dim'],
self.arch_params['n_hidden_0'], tf.nn.relu,
'dense0', reuse)
h1 = ops.dense(h0, self.arch_params['n_hidden_0'],
self.arch_params['n_hidden_1'], tf.nn.relu,
'dense1', reuse)
relu1_do = tf.nn.dropout(h1, self.arch_params['do_keep_prob'])
d = ops.dense(relu1_do, self.arch_params['n_hidden_1'],
self.arch_params['out_dim'], None, 'dense2', reuse)
return d
def get_agent(x, reuse=False):
"""
Generate the CNN agent
:param x: tensor, Input frames concatenated along axis 3
:param reuse: bool, True -> Reuse weight variables
False -> Create new ones
:return: Tensor, logits for each valid action
"""
if reuse:
tf.get_variable_scope().reuse_variables()
x = tf.reshape(x, shape=[-1, 128*4])
dense_1 = tf.nn.relu(ops.dense(x, 128*4, mc.dense_1, name='dense_1'))
dense_2 = tf.nn.relu(ops.dense(dense_1, mc.dense_1, mc.dense_2, name='dense_2'))
dense_3 = tf.nn.relu(ops.dense(dense_2, mc.dense_2, mc.dense_3, name='dense_3'))
output = ops.dense(dense_3, mc.dense_3, env.action_space.n, name='output')
return output
def discriminator(x, reuse=False):
if reuse:
tf.get_variable_scope().reuse_variables()
conv_b_1 = ops.conv_block(x,
filter_size=4,
stride_length=2,
n_maps=8,
name='d_conv_b_1')
conv_b_2 = ops.conv_block(conv_b_1,
filter_size=4,
stride_length=2,
n_maps=16,
name='d_conv_b_2')
conv_b_3 = ops.conv_block(conv_b_2,
filter_size=4,
stride_length=2,
n_maps=32,
name='d_conv_b_3')
conv_b_4 = ops.conv_block(conv_b_3,
filter_size=4,
stride_length=2,
n_maps=64,
name='d_conv_b_4')
conv_b_5 = ops.conv_block(conv_b_4,
filter_size=4,
stride_length=2,
n_maps=1,
name='d_conv_b_5')
conv_b_5_r = tf.reshape(conv_b_5, [-1, 11 * 9 * 1], name='d_reshape')
output = ops.dense(conv_b_5_r, 11 * 9, 1, name='d_output')
return output
def discriminator(x, reuse=False):
if reuse:
tf.get_variable_scope().reuse_variables()
conv_1 = ops.lrelu(
ops.batch_norm(ops.cnn_2d(x,
weight_shape=[4, 4, 6, 64],
strides=[1, 2, 2, 1],
name='dis_conv_1'),
center=True,
scale=True,
is_training=True,
scope='dis_batch_Norm_1'))
conv_2 = ops.lrelu(
ops.batch_norm(ops.cnn_2d(conv_1,
weight_shape=[4, 4, 64, 128],
strides=[1, 2, 2, 1],
name='dis_conv_2'),
center=True,
scale=True,
is_training=True,
scope='dis_batch_Norm_2'))
conv_3 = ops.lrelu(
ops.batch_norm(ops.cnn_2d(conv_2,
weight_shape=[4, 4, 128, 256],
strides=[1, 2, 2, 1],
name='dis_conv_3'),
center=True,
scale=True,
is_training=True,
scope='dis_batch_Norm_3'))
conv_4 = ops.lrelu(
ops.batch_norm(ops.cnn_2d(conv_3,
weight_shape=[4, 4, 256, 512],
strides=[1, 2, 2, 1],
name='dis_conv_4'),
center=True,
scale=True,
is_training=True,
scope='dis_batch_Norm_4'))
conv_5 = ops.lrelu(
ops.batch_norm(ops.cnn_2d(conv_4,
weight_shape=[4, 4, 512, 512],
strides=[1, 2, 2, 1],
name='dis_conv_5'),
center=True,
scale=True,
is_training=True,
scope='dis_batch_Norm_5'))
conv_6 = ops.lrelu(
ops.batch_norm(ops.cnn_2d(conv_5,
weight_shape=[4, 4, 512, 512],
strides=[1, 2, 2, 1],
name='dis_conv_6'),
center=True,
scale=True,
is_training=True,
scope='dis_batch_Norm_6'))
output = ops.dense(conv_6, 5 * 6, 1, name='dis_output')
return output
def _build_decoder(self, z):
'''
Deconvolutional decoder; takes a sampled posterior, reshapes
and deconvolves it using shared encoder weights.
'''
kernels = self._kernels
kernels.reverse()
# Project and reshape the sampled latent variable so that it can be convolved
h, w, d = self.encoding.get_shape().as_list()[1:]
flat_dim = h * w * d
layer_inputs = dense(flat_dim, z)
layer_inputs = tf.reshape(layer_inputs, [-1, h, w, d])
# Encoding convolutions are applied in reverse
for i in range(len(kernels)):
depth_out = layer_inputs.get_shape().as_list()[-1]
W = kernels[i]
b = tf.Variable(tf.zeros([depth_out]))
# Deconvolution with shared encoder weights
layer_inputs = tf.nn.conv2d_transpose(input=layer_inputs,
filter=W,
strides=[1, 1, 1, 1],
padding='SAME',
name='dec_conv{}'.format(i))
layer_inputs = tf.add(layer_inputs, b)
layer_inputs = tf.nn.relu(layer_inputs)
kernels.reverse()
return layer_inputs
def regression_loss(self, hidden, labels, is_training, scope,
reuse=tf.AUTO_REUSE, return_logits=False):
"""Get regression loss."""
net_config = self.net_config
initializer = self.get_initializer()
with tf.variable_scope(scope, reuse=reuse):
hidden = ops.dropout_op(hidden, net_config.dropout, training=is_training)
logits = ops.dense(
hidden,
1,
initializer=initializer,
scope="logit")
# Always cast to float32 for loss
logits = tf.squeeze(logits, axis=-1)
if logits.dtype != tf.float32:
logits = tf.cast(logits, tf.float32)
loss = tf.square(logits - tf.cast(labels, logits.dtype))
if return_logits:
return loss, logits
return loss
def discriminator(self, inputs, scope='discriminator', reuse=None):
with tf.variable_scope(scope, reuse=reuse):
# Set discriminator to always be training. Reason for doing this is
# For the WGAN gradient loss (which is not the default loss function for
# this model, still uses this architecture), the loss function has an expression
# which is the gradient of an instance of the discriminator. Putting that
# into the optimizer creates a dependency on the second order gradient of the
# disriminator. Batch normalization where the training vs running flag is itself
# a TF variable (rather than normal python boolean) seems to break this. Easier to
# just set to True because in this model we only ever use the discriminator for
# training (to train the generator).
bn = BN(True)
t = lrelu(conv2d(inputs, 64)) # no bn here
t = lrelu(bn(conv2d(t, 128)))
t = lrelu(bn(conv2d(t, 256)))
t = lrelu(bn(conv2d(t, 512)))
t = lrelu(bn(conv2d(t, 1024)))
# flatten 3D tensor into 1D to prepare for a dense (fully connected)
# layer. Flattened tensor can also be treated as vector that can be
# used for learned similarty measurements between images.
similarity = flatten(t)
# return logits (before sigmoid activation) because several TF
# accumulator functions expect logits, and do the sigmoid for you
logits = dense(similarity, 1)
classification = sigmoid(logits)
return classification, logits, similarity
def generator(self, inputs, scope='generator', reuse=None):
with tf.variable_scope(scope, reuse=reuse):
# generator will upscale a tiny image with layers of convolution
# until it reaches the final output image dimensions. These vars
# are the starting tiny image dimensions.
# This is why this arch needs images that are divisible by 32
minirows = self.img_shape[0] // 32
minicols = self.img_shape[1] // 32
# batch normalization, which needs to know whether this is training or
# application
bn = BN(self.is_training)
# dense (i.e. fully connected) layer followed by reshaping into the tiny
# image. The tiny image has a Z dim of 512 that gradually gets reduced
# to 3 channels (r, g, b)
t = dense(inputs, minirows * minicols * 512)
t = lrelu(bn(reshape(t,
(tf.shape(t)[0], minirows, minicols, 512))))
t = lrelu(bn(conv2dtr(t, 512)))
t = lrelu(bn(conv2dtr(t, 256)))
t = lrelu(bn(conv2dtr(t, 128)))
t = lrelu(bn(conv2dtr(t, 64)))
# final conv2d transpose to get to filter depth of 3, for rgb channels
logits = conv2dtr(t, self.img_shape[2])
return tanh(logits) # common final activation in GANs
def one_hidden_mlp(inpt, units, outsize):
"""
Fully connected model, with one hidden layer of length units and relu
nonlinearity.
Args:
- inpt (Tensor) : input of the model;
- units (int) : numberr of hidden units.
Returns:
- output (Tensor) : output of the model.
"""
with tf.variable_scope("one_hidden"):
hidden = tf.nn.relu(dense(inpt, units=units, name="dense1"))
output = dense(hidden, units=outsize, name="dense2")
return output
def classification_loss(self, hidden, labels, n_class, is_training, scope,
reuse=tf.AUTO_REUSE, return_logits=False):
"""Get classification loss."""
net_config = self.net_config
initializer = self.get_initializer()
with tf.variable_scope(scope, reuse=reuse):
hidden = ops.dropout_op(hidden, net_config.dropout, training=is_training)
logits = ops.dense(
hidden,
n_class,
initializer=initializer,
scope="logit")
# Always cast to float32 for softmax & loss
if logits.dtype != tf.float32:
logits = tf.cast(logits, tf.float32)
one_hot_target = tf.one_hot(labels, n_class, dtype=logits.dtype)
loss = -tf.reduce_sum(tf.nn.log_softmax(logits) * one_hot_target, -1)
if return_logits:
return loss, logits
return loss
def dec(x, start_res, end_res, scope='Decoder'):
with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
x = ops.dense('fc1', x, fn(4) * 4 * 4, 'NHWC')
x = tf.reshape(x, [-1, 4, 4, fn(4)])
res = 8
prev_x = None
while res <= end_res:
prev_x = x
x = block_up(x, fn(res), 3, rname(res))
res *= 2
res = res // 2
if res > start_res:
t = tf.get_variable(
rname(res) + '_t',
shape=[],
collections=[tf.GraphKeys.GLOBAL_VARIABLES, "lerp"],
dtype=tf.float32,
initializer=tf.zeros_initializer(),
trainable=False)
x1 = ops.to_rgb('rgb_' + rname(res // 2), prev_x, 'NHWC')
x1 = ops.upscale2d(x1, 'NHWC')
x2 = ops.to_rgb('rgb_' + rname(res), x, 'NHWC')
x = ops.lerp_clip(x1, x2, t)
else:
x = ops.to_rgb('rgb_' + rname(res), x, "NHWC")
x_shape = utils.int_shape(x)
assert (end_res == x_shape[1])
assert (end_res == x_shape[2])
return x
def enc(x, start_res, end_res, scope='Encoder'):
with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
res = end_res
if res > start_res:
x1 = ops.downscale2d(x, 'NHWC')
x1 = ops.from_rgb('rgb_' + rname(res // 2), x1, fn(res // 2),
'NHWC')
x2 = ops.from_rgb('rgb_' + rname(res), x, fn(res // 2), 'NHWC')
t = tf.get_variable(
rname(res) + '_t',
shape=[],
dtype=tf.float32,
collections=[tf.GraphKeys.GLOBAL_VARIABLES, "lerp"],
initializer=tf.zeros_initializer(),
trainable=False)
x2 = block_dn(x2, fn(res), fn(res // 2), 3, rname(res))
x = ops.lerp_clip(x1, x2, t)
res = res // 2
else:
x = ops.from_rgb('rgb_' + rname(res), x, fn(res), 'NHWC')
while res >= 4:
x = block_dn(x, fn(res), fn(res // 2), 3, rname(res))
res = res // 2
x = tf.layers.flatten(x)
x = ops.dense('fc1', x, 512, 'NHWC')
mean, std = tf.split(x, 2, 1)
return mean, std
def __init__(self, sess, n_classes, img_height, img_width, lr_init,
lr_decay, lr_step, lr_min, model_dir, mode):
self.n_classes = n_classes
self.model_dir = model_dir
self.lr_init = lr_init
self.lr_decay = lr_decay
self.lr_step = lr_step
self.lr_min = lr_min
# Input params
self.input = tf.placeholder(
'float32', [None, img_height, img_width, 1]) # [batch_size, ...]
self.labels = tf.placeholder('float32', [None, n_classes])
self.keep_prob = tf.placeholder('float32', name='keep_prob')
self.global_step = tf.Variable(0, trainable=False)
if mode == 'inception':
self.net = inception_cnn(self.input, self.keep_prob)
elif mode == 'resnet':
self.net = resnet_cnn(self.input, self.keep_prob)
else:
self.net = vgg_cnn(self.input, self.keep_prob)
self.flat = tf.contrib.layers.flatten(self.net)
#self.dense, self.dense_w, self.dense_b = dense(self.flat, tf.nn.relu, 100)
self.logits, self.out_w, self.out_b = dense(self.flat,
None,
self.n_classes,
name='out')
# Define loss
self.losses = tf.nn.softmax_cross_entropy_with_logits_v2(
logits=self.logits, labels=self.labels)
# compute mean loss
self.loss = tf.reduce_mean(self.losses)
# run optimization
self.learning_rate_op = tf.maximum(
self.lr_min,
tf.train.exponential_decay(self.lr_init,
self.global_step,
self.lr_step,
self.lr_decay,
staircase=True))
self.train_op = tf.train.AdamOptimizer(self.learning_rate_op).minimize(
self.loss)
# Evaluate model
self.predictions = tf.map_fn(tf.nn.softmax, self.logits)
correct_pred = tf.equal(tf.argmax(self.predictions, axis=1),
tf.argmax(self.labels, axis=1))
self.accuracy = tf.reduce_mean(tf.cast(correct_pred, tf.float32))
self.constructSummary(sess)
def discriminator(x, resolution, cfg, is_training=True, scope='Discriminator'):
assert (cfg.data_format == 'NCHW' or cfg.data_format == 'NHWC')
def rname(resolution):
return str(resolution) + 'x' + str(resolution)
def fmap(resolution):
return cfg.resolution_to_filt_num[resolution]
x_shape = utils.int_shape(x)
assert (resolution == x_shape[1 if cfg.data_format == 'NHWC' else 3])
assert (resolution == x_shape[2])
with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):
if resolution > cfg.starting_resolution:
x1 = ops.downscale2d(x, cfg.data_format)
x1 = ops.from_rgb('from_rgb_' + rname(resolution // 2), x1,
fmap(resolution // 2), cfg.data_format)
x2 = ops.from_rgb('from_rgb_' + rname(resolution), x,
fmap(resolution // 2), cfg.data_format)
t = tf.get_variable(
rname(resolution) + '_t',
shape=[],
dtype=tf.float32,
collections=[tf.GraphKeys.GLOBAL_VARIABLES, "lerp"],
initializer=tf.zeros_initializer(),
trainable=False)
num_filters = [fmap(resolution), fmap(resolution // 2)]
x2 = dblock(rname(resolution), x2, num_filters, cfg.data_format)
x = ops.lerp_clip(x1, x2, t)
resolution = resolution // 2
else:
x = ops.from_rgb('from_rgb_' + rname(resolution), x,
fmap(resolution), cfg.data_format)
while resolution >= 4:
if resolution == 4:
x = ops.minibatch_stddev_layer(x, cfg.data_format)
num_filters = [fmap(resolution), fmap(resolution // 2)]
x = dblock(rname(resolution), x, num_filters, cfg.data_format)
resolution = resolution // 2
x = ops.dense('2x2', x, fmap(resolution), cfg.data_format)
x = ops.leaky_relu(x)
x = ops.dense('output', x, 1, cfg.data_format)
return x
def generator(self, x, c, reuse=False):
print("Generator ...........")
with tf.variable_scope('generator') as scope:
if (reuse):
scope.reuse_variables()
inputs = tf.concat([x, c], axis=-1)
# Dense layer 1
dense1 = ops.dense(inputs,
64 * 8 * 8,
0.02,
self.training,
norm=self.g_norm,
scope="g_inputs")
dense1 = tf.nn.relu(dense1)
dense1 = tf.reshape(dense1, [-1, 8, 8, 64])
outputs = dense1
# Res Block
for i in range(self.res_block_size):
outputs = resblock1_G(outputs,
64,
3,
1,
self.training,
norm=self.g_norm,
scope='g_residual_{:d}'.format(i))
outputs = ops.batch_norm(outputs, self.training)
outputs = tf.nn.relu(outputs)
outputs = outputs + dense1
# Upscaling by pixel shuffling
for i in range(3):
outputs = resblock2_G(outputs,
256,
3,
1,
2,
self.training,
norm=None,
scope='g_upscale_{:d}'.format(i))
outputs = ops.conv_2d(outputs,
3,
9,
1,
padding='SAME',
stddev=0.02,
training=self.training,
norm=None,
scope="g_conv_last")
outputs = tf.nn.tanh(outputs)
return outputs
def predict(pfm, ufm, is_training, keep_prob):
fm = tf.concat([pfm, ufm], axis=3) # 7*7*512
network = conv2d(fm, [2, 2, 512, 1024], [1, 1, 1, 1],
is_training,
bn=True,
scope='conv1_1')
network = tf.nn.max_pool(network,
ksize=[1, 2, 2, 1],
strides=[1, 2, 2, 1],
padding='VALID',
name='pool1') # [3, 3, 1024]
network = conv2d(network, [2, 2, 1024, 2048], [1, 1, 1, 1],
is_training,
bn=True,
scope='conv2_1')
network = tf.nn.max_pool(network,
ksize=[1, 2, 2, 1],
strides=[1, 2, 2, 1],
padding='VALID',
name='pool2') # [32, 1, 1, 2048]
fea_vec = tf.reshape(network, [-1, 2048])
network = tf.nn.dropout(
dense(fea_vec, 2048, 256, is_training, bn=True, scope='mlp1'),
keep_prob)
network = tf.nn.dropout(
dense(network, 256, 128, is_training, bn=True, scope='mlp2'),
keep_prob)
network = tf.nn.dropout(
dense(network, 128, 64, is_training, bn=True, scope='mlp3'), keep_prob)
network = tf.nn.dropout(
dense(network,
64,
1,
is_training,
activation_fn=tf.nn.tanh,
bn=True,
scope='mlp4'), keep_prob)
return tf.reshape(network, [-1])
def build_model(self, inputs, labels, is_training=False):
self.network = ops.convolution(inputs,
self.channels,
50,
5,
50,
is_training=is_training,
scope='conv1')
self.network = ops.pooling(self.network, scope='pool1')
self.network = ops.convolution(self.network,
50,
20,
5,
20,
is_training=is_training,
scope='conv2')
self.network = ops.pooling(self.network, scope='pool2')
self.network = ops.flatten(self.network, scope='flatten')
self.network = ops.dense(self.network,
self.network.get_shape().as_list()[1],
200,
scope='fc1')
self.network = ops.dense(self.network, 200, 50, scope='fc2')
self.network = ops.dense(self.network,
50,
10,
activation=None,
scope='fc3')
self.loss = ops.loss(self.network, labels, scope='loss')
self.accuracy = ops.accuracy(self.network, labels, scope='accuracy')
if is_training:
self.optimizer = ops.optimize(self.loss,
self.learning_rate,
scope='update')
def __call__(self, x, n_channels, n_codes):
non_linearity = tf.nn.relu
with tf.variable_scope(self.name, reuse=self.reuse):
conv1_a = conv2d(x, 'conv1_a', 16, 3, 2, 'SAME', True,
non_linearity, self.is_train)
conv1_b = conv2d(conv1_a, 'conv1_b', 16, 3, 1, 'SAME', True,
non_linearity, self.is_train)
conv2_a = conv2d(conv1_b, 'conv2_a', 32, 3, 2, 'SAME', True,
non_linearity, self.is_train)
conv2_b = conv2d(conv2_a, "conv2_b", 32, 3, 1, 'SAME', True,
non_linearity, self.is_train)
conv3 = conv2d(conv2_b, 'conv3', 1, 3, 2, 'SAME', True,
non_linearity, self.is_train)
fc1 = dense(tf.reshape(conv3, [-1, np.prod([8, 8, 1])]), 'fc1',
n_codes, False, None, self.is_train)
fc2 = dense(fc1, 'fc2', np.prod([8, 8, 1]), False, non_linearity,
self.is_train)
deconv4 = upconv2d(tf.reshape(fc2, [-1] + [8, 8, 1]), 2, 'deconv4',
32, 3, 1, 'SAME', True, non_linearity,
self.is_train)
conv4 = conv2d(deconv4, 'conv4', 32, 3, 1, 'SAME', True,
non_linearity, self.is_train)
deconv5 = upconv2d(conv4, 2, 'deconv5', 16, 3, 1, 'SAME', True,
non_linearity, self.is_train)
conv5 = conv2d(deconv5, 'conv5', 16, 3, 1, 'SAME', True,
non_linearity, self.is_train)
deconv6 = upconv2d(conv5, 2, 'deconv6', 16, 3, 1, 'SAME', True,
non_linearity, self.is_train)
conv_out = conv2d(deconv6, 'conv_out', n_channels, 3, 1, 'SAME',
False, None, self.is_train)
if self.reuse is None:
self.var_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES,
scope=self.name)
self.saver = tf.train.Saver(var_list=self.var_list,
max_to_keep=100)
self.reuse = True
return fc1, conv_out
