def define_graph(self):
batch_size = self._batch_size
self.input_frame = tf.placeholder(
tf.float32, shape=[batch_size, 31, 31, 10], name="s1"
)
self.input_frame_2 = tf.placeholder(
tf.float32, shape=[batch_size, 23, 23, 10], name="s2"
)
self.input_frame_3 = tf.placeholder(
tf.float32, shape=[batch_size, 11, 11, 10], name="s3"
)
self.rain_train = tf.placeholder(
tf.float32, shape=[batch_size, 1], name="rain"
)
self.rain_train_2 = tf.placeholder(
tf.float32, shape=[batch_size, 1], name="rain_2"
)
mid_27 = conv2d(self.input_frame, w([5, 5, 10, 32]), b([32, ]))
print(mid_27)
mid_23 =conv2d(mid_27, w([5, 5, 32, 64]), b([64, ]))
print(mid_23)
mid_23_2 = tf.concat([mid_23, self.input_frame_2], 3)
mid_19 = conv2d(mid_23_2, w([5, 5, 64 + 10, 128]), b([128, ]))
mid_15 = conv2d(mid_19, w([5, 5, 128, 256]), b([256, ]))
mid_11 = conv2d(mid_15, w([5, 5, 256, 256]), b([256, ]))
print(mid_11)
mid_11_2 = tf.concat([mid_11, self.input_frame_3], 3)
mid_9 = conv2d(mid_11_2, w([3, 3, 256 + 10, 128]), b([128, ]))
mid_7 = conv2d(mid_9, w([3, 3, 128, 64]), b([64, ]))
mid_5 = conv2d(mid_7, w([3, 3, 64, 32]), b([32, ]))
mid_3 = conv2d(mid_5, w([3, 3, 32, 16]), b([16, ]))
regress = conv2d(mid_3, w([3, 3, 16, 1]), b([1, ]))
regress = tf.reshape(regress, (batch_size, 1))
self.loss = tf.losses.mean_squared_error(regress, self.rain_train)
self.bias = tf.reduce_mean(tf.abs(regress - self.rain_train))
self.error_rate = tf.reduce_mean(tf.abs(regress - self.rain_train) / self.rain_train_2) * 100
self.optimizer = tf.train.AdamOptimizer(self._learning_rate).minimize(self.loss)
self.pred_test = regress
def define_graph(self):
# Sets up the model graph
# The variables to train
self.train_vars = []
# Sets up the layer
with tf.name_scope('net'):
with tf.name_scope('setup'):
# Convolution
with tf.name_scope('convolutions'):
self.conv_ws = []
self.conv_bs = []
last_out_height = self.height
last_out_width = self.width
# last_out_seqlen = self.seqlen
for i in xrange(len(self.kernel_sizes)):
self.conv_ws.append(
w([
self.kernel_sizes[i], self.kernel_sizes[i],
self.feature_maps[i], self.feature_maps[i + 1]
]))
self.conv_bs.append(b([self.feature_maps[i + 1]]))
last_out_height = conv_out_size(
last_out_height, 'SAME', self.kernel_sizes[i], 1)
last_out_width = conv_out_size(last_out_width, 'SAME',
self.kernel_sizes[i], 1)
with tf.name_scope('fully-connected'):
# Add in an initial layer to go from the last conv to the first fully-connected.
# Use /2 for the height and width because there is a 2x2 pooling layer
self.fc_layer_sizes.insert(0, (last_out_height / 2) *
(last_out_width / 2) *
self.feature_maps[-1])
self.fc_ws = []
self.fc_bs = []
for i in xrange(len(self.fc_layer_sizes) - 1):
self.fc_ws.append(
w([
self.fc_layer_sizes[i],
self.fc_layer_sizes[i + 1]
]))
self.fc_bs.append(b([self.fc_layer_sizes[i + 1]]))
self.train_vars += self.conv_ws
self.train_vars += self.conv_bs
self.train_vars += self.fc_ws
self.train_vars += self.fc_bs
def define_graph(self):
batch_size = self._batch_size
self.input_frame = tf.placeholder(
tf.float32,
shape=[batch_size, self.frame_size, self.frame_size, 1],
name="s1")
self.input_frame_2 = tf.placeholder(tf.float32,
shape=[batch_size, 27, 27, 1],
name="s2")
self.input_frame_3 = tf.placeholder(tf.float32,
shape=[batch_size, 11, 11, 1],
name="s3")
self.rain_train = tf.placeholder(tf.float32,
shape=[batch_size, 1],
name="rain")
mid1 = conv2d(self.input_frame, w([8, 8, 1, 32]), b([
32,
]))
mid2 = max_pool2d(mid1)
print(mid2)
mid3 = tf.concat([mid2, self.input_frame_2], 3)
print(mid3)
mid4 = conv2d(mid3, w([6, 6, 32 + 1, 64]), b([
64,
]))
mid5 = max_pool2d(mid4)
mid6 = tf.concat([mid5, self.input_frame_3], 3)
mid7 = conv2d(mid6, w([3, 3, 64 + 1, 64]), b([
64,
]))
print(mid7)
mid8 = tf.reshape(mid7, [-1, 9 * 9 * 64])
dropout = tf.nn.dropout(mid8, keep_prob=0.5)
regress = tf.nn.xw_plus_b(dropout, w([9 * 9 * 64, 1]), b([1]))
self.loss = tf.losses.mean_squared_error(regress, self.rain_train)
self.bias = tf.reduce_mean(tf.abs(regress - self.rain_train))
self.optimizer = tf.train.AdamOptimizer(self._learning_rate).minimize(
self.loss)
def define_variables(self, layer_variables=None):
"Define trainable variables"
if self.verbose: print("Defining Wavenet layer variables..")
if layer_variables == None:
with tf.name_scope("wavenet_params"):
## DEFINE VARIABLES
in_channels = self.in_channels
layer_weights, layer_biases = [], []
for block_num in range(self.num_blocks): # FOR EACH BLOCK
for layer_num in range(self.num_layers): # FOR EACH LAYER
name = "block_%i-layer_%i" % (block_num, layer_num)
with tf.name_scope(name):
# WEIGHT INITIALISATION (==sqrt(k / number of inputs per neuron)), k=2 for relu, k=1 for tanh
if self.activation == tf.nn.relu:
stddev = np.sqrt(2) / np.sqrt(
self.filter_width * in_channels)
elif self.activation == tf.nn.tanh:
stddev = np.sqrt(1) / np.sqrt(
self.filter_width * in_channels)
else:
stddev = 0.01
weights = w(shape=(self.filter_width, in_channels,
self.hidden_channels),
mean=0.,
stddev=stddev,
name="weights")
#weights = w(shape=(self.filter_width, in_channels, self.hidden_channels), mean=1, stddev=0, name="weights")# for testing purposes only
layer_weights.append(weights)
if self.biases:
if self.activation == tf.nn.relu: const = 0.1
else: const = 0.0
biases = b(shape=(1, 1, self.hidden_channels),
const=const,
name="biases")
layer_biases.append(biases)
in_channels = self.hidden_channels
else:
layer_weights = layer_variables[0]
layer_biases = layer_variables[1]
self.layer_weights = layer_weights
self.layer_biases = layer_biases
self.num_weights = np.sum(
[weights.shape.num_elements() for weights in self.layer_weights])
self.num_biases = np.sum(
[biases.shape.num_elements() for biases in self.layer_biases])
return (layer_weights, layer_biases)
def define_graph(self):
"""
Sets up the model graph in TensorFlow.
"""
with tf.name_scope('generator'):
##
# Data
##
with tf.name_scope('data'):
self.input_frames_train = tf.placeholder(
tf.float32, shape=[None, self.height_train, self.width_train, 3 * c.HIST_LEN])
self.gt_frames_train = tf.placeholder(
tf.float32, shape=[None, self.height_train, self.width_train, 3])
self.input_frames_test = tf.placeholder(
tf.float32, shape=[None, self.height_test, self.width_test, 3 * c.HIST_LEN])
self.gt_frames_test = tf.placeholder(
tf.float32, shape=[None, self.height_test, self.width_test, 3])
# use variable batch_size for more flexibility
self.batch_size_train = tf.shape(self.input_frames_train)[0]
self.batch_size_test = tf.shape(self.input_frames_test)[0]
##
# Scale network setup and calculation
##
self.summaries_train = []
self.scale_preds_train = [] # the generated images at each scale
self.scale_gts_train = [] # the ground truth images at each scale
self.d_scale_preds = [] # the predictions from the discriminator model
self.summaries_test = []
self.scale_preds_test = [] # the generated images at each scale
self.scale_gts_test = [] # the ground truth images at each scale
for scale_num in range(self.num_scale_nets):
with tf.name_scope('scale_' + str(scale_num)):
with tf.name_scope('setup'):
ws = []
bs = []
# create weights for kernels
for i in range(len(self.scale_kernel_sizes[scale_num])):
ws.append(w([self.scale_kernel_sizes[scale_num][i],
self.scale_kernel_sizes[scale_num][i],
self.scale_layer_fms[scale_num][i],
self.scale_layer_fms[scale_num][i + 1]]))
bs.append(b([self.scale_layer_fms[scale_num][i + 1]]))
with tf.name_scope('calculation'):
def calculate(height, width, inputs, gts, last_gen_frames):
# scale inputs and gts
scale_factor = 1. / 2 ** ((self.num_scale_nets - 1) - scale_num)
scale_height = int(height * scale_factor)
scale_width = int(width * scale_factor)
inputs = tf.image.resize_images(inputs, [scale_height, scale_width])
scale_gts = tf.image.resize_images(gts, [scale_height, scale_width])
# for all scales but the first, add the frame generated by the last
# scale to the input
if scale_num > 0:
last_gen_frames = tf.image.resize_images(
last_gen_frames,[scale_height, scale_width])
print("inputs: {}, frames: {}".format(inputs.shape, last_gen_frames.shape))
inputs = tf.concat([inputs, last_gen_frames], 3)
# generated frame predictions
preds = inputs
# perform convolutions
with tf.name_scope('convolutions'):
for i in range(len(self.scale_kernel_sizes[scale_num])):
# Convolve layer
preds = tf.nn.conv2d(
preds, ws[i], [1, 1, 1, 1], padding=c.PADDING_G)
# Activate with ReLU (or Tanh for last layer)
if i == len(self.scale_kernel_sizes[scale_num]) - 1:
preds = tf.nn.tanh(preds + bs[i])
else:
preds = tf.nn.relu(preds + bs[i])
return preds, scale_gts
##
# Perform train calculation
##
# for all scales but the first, add the frame generated by the last
# scale to the input
if scale_num > 0:
last_scale_pred_train = self.scale_preds_train[scale_num - 1]
else:
last_scale_pred_train = None
# calculate
train_preds, train_gts = calculate(self.height_train,
self.width_train,
self.input_frames_train,
self.gt_frames_train,
last_scale_pred_train)
self.scale_preds_train.append(train_preds)
self.scale_gts_train.append(train_gts)
# We need to run the network first to get generated frames, run the
# discriminator on those frames to get d_scale_preds, then run this
# again for the loss optimization.
if c.ADVERSARIAL:
self.d_scale_preds.append(tf.placeholder(tf.float32, [None, 1]))
##
# Perform test calculation
##
# for all scales but the first, add the frame generated by the last
# scale to the input
if scale_num > 0:
last_scale_pred_test = self.scale_preds_test[scale_num - 1]
else:
last_scale_pred_test = None
# calculate
test_preds, test_gts = calculate(self.height_test,
self.width_test,
self.input_frames_test,
self.gt_frames_test,
last_scale_pred_test)
self.scale_preds_test.append(test_preds)
self.scale_gts_test.append(test_gts)
##
# Training
##
with tf.name_scope('train'):
# global loss is the combined loss from every scale network
self.global_loss = combined_loss(self.scale_preds_train,
self.scale_gts_train,
self.d_scale_preds)
self.global_step = tf.Variable(0, trainable=False)
self.optimizer = tf.train.AdamOptimizer(learning_rate=c.LRATE_G, name='optimizer')
self.train_op = self.optimizer.minimize(self.global_loss,
global_step=self.global_step,
name='train_op')
# train loss summary
loss_summary = tf.summary.scalar('train_loss_G', self.global_loss)
self.summaries_train.append(loss_summary)
##
# Error
##
with tf.name_scope('error'):
# error computation
# get error at largest scale
self.psnr_error_train = psnr_error(self.scale_preds_train[-1],
self.gt_frames_train)
self.sharpdiff_error_train = sharp_diff_error(self.scale_preds_train[-1],
self.gt_frames_train)
self.psnr_error_test = psnr_error(self.scale_preds_test[-1],
self.gt_frames_test)
self.sharpdiff_error_test = sharp_diff_error(self.scale_preds_test[-1],
self.gt_frames_test)
# train error summaries
summary_psnr_train = tf.summary.scalar('train_PSNR',
self.psnr_error_train)
summary_sharpdiff_train = tf.summary.scalar('train_SharpDiff',
self.sharpdiff_error_train)
self.summaries_train += [summary_psnr_train, summary_sharpdiff_train]
# test error
summary_psnr_test = tf.summary.scalar('test_PSNR',
self.psnr_error_test)
summary_sharpdiff_test = tf.summary.scalar('test_SharpDiff',
self.sharpdiff_error_test)
self.summaries_test += [summary_psnr_test, summary_sharpdiff_test]
# add summaries to visualize in TensorBoard
self.summaries_train = tf.summary.merge(self.summaries_train)
self.summaries_test = tf.summary.merge(self.summaries_test)
def define_graph(self):
"""
Sets up the model graph in TensorFlow.
"""
##
# Input data
##
with tf.name_scope('input'):
self.input_frames = tf.placeholder(
tf.float32, shape=[None, self.height, self.width, self.conv_layer_fms[0]])
# use variable batch_size for more flexibility
self.batch_size = tf.shape(self.input_frames)[0]
##
# Layer setup
##
with tf.name_scope('setup'):
# convolution
with tf.name_scope('convolutions'):
conv_ws = []
conv_bs = []
last_out_height = self.height
last_out_width = self.width
for i in range(len(self.kernel_sizes)):
conv_ws.append(w([self.kernel_sizes[i],
self.kernel_sizes[i],
self.conv_layer_fms[i],
self.conv_layer_fms[i + 1]]))
conv_bs.append(b([self.conv_layer_fms[i + 1]]))
last_out_height = conv_out_size(
last_out_height, c.PADDING_D, self.kernel_sizes[i], 1)
last_out_width = conv_out_size(
last_out_width, c.PADDING_D, self.kernel_sizes[i], 1)
# fully-connected
with tf.name_scope('full-connected'):
# Add in an initial layer to go from the last conv to the first fully-connected.
# Use /2 for the height and width because there is a 2x2 pooling layer
self.fc_layer_sizes.insert(
0, (last_out_height / 2) * (last_out_width / 2) * self.conv_layer_fms[-1])
fc_ws = []
fc_bs = []
for i in range(len(self.fc_layer_sizes) - 1):
fc_ws.append(w([int(self.fc_layer_sizes[i]),
int(self.fc_layer_sizes[i + 1])]))
fc_bs.append(b([int(self.fc_layer_sizes[i + 1])]))
##
# Forward pass calculation
##
def generate_predictions():
"""
Runs self.input_frames through the network to generate a prediction from 0
(generated img) to 1 (real img).
@return: A tensor of predictions of shape [self.batch_size x 1].
"""
with tf.name_scope('calculation'):
preds = tf.zeros([self.batch_size, 1])
last_input = self.input_frames
# convolutions
with tf.name_scope('convolutions'):
for i in range(len(conv_ws)):
# Convolve layer and activate with ReLU
preds = tf.nn.conv2d(
last_input, conv_ws[i], [1, 1, 1, 1], padding=c.PADDING_D)
preds = tf.nn.relu(preds + conv_bs[i])
last_input = preds
# pooling layer
with tf.name_scope('pooling'):
preds = tf.nn.max_pool(preds, [1, 2, 2, 1], [1, 2, 2, 1], padding=c.PADDING_D)
# flatten preds for dense layers
shape = preds.get_shape().as_list()
# -1 can be used as one dimension to size dynamically
preds = tf.reshape(preds, [-1, shape[1] * shape[2] * shape[3]])
# fully-connected layers
with tf.name_scope('fully-connected'):
for i in range(len(fc_ws)):
preds = tf.matmul(preds, fc_ws[i]) + fc_bs[i]
# Activate with ReLU (or Sigmoid for last layer)
if i == len(fc_ws) - 1:
preds = tf.sigmoid(preds)
else:
preds = tf.nn.relu(preds)
# clip preds between [.1, 0.9] for stability
with tf.name_scope('clip'):
preds = tf.clip_by_value(preds, 0.1, 0.9)
return preds
self.preds = generate_predictions()
def define_graph(self):
"""
Sets up the model graph in TensorFlow.
"""
##
# Input data
##
with tf.name_scope('input'):
self.input_frames = tf.placeholder(
tf.float32,
shape=[None, self.height, self.width, self.conv_layer_fms[0]])
# use variable batch_size for more flexibility
self.batch_size = tf.shape(self.input_frames)[0] / (c.HIST_LEN + 1)
# self.batch_size = c.BATCH_SIZE
##
# Layer setup
##
with tf.name_scope('setup'):
# convolution
with tf.name_scope('convolutions'):
conv_ws = []
conv_bs = []
last_out_height = self.height
last_out_width = self.width
for i in xrange(len(self.kernel_sizes)):
conv_ws.append(
w([
self.kernel_sizes[i], self.kernel_sizes[i],
self.conv_layer_fms[i], self.conv_layer_fms[i + 1]
]))
conv_bs.append(b([self.conv_layer_fms[i + 1]]))
last_out_height = conv_out_size(last_out_height,
c.PADDING_D,
self.kernel_sizes[i], 1)
last_out_width = conv_out_size(last_out_width, c.PADDING_D,
self.kernel_sizes[i], 1)
# lstm
with tf.name_scope('lstm'):
hidden_dim = (last_out_height / 2) * (
last_out_width / 2) * self.conv_layer_fms[-1]
# fully-connected
with tf.name_scope('full-connected'):
# Add in an initial layer to go from the last conv to the first fully-connected.
# Use /2 for the height and width because there is a 2x2 pooling layer
self.fc_layer_sizes.insert(0, hidden_dim)
fc_ws = []
fc_bs = []
for i in xrange(len(self.fc_layer_sizes) - 1):
fc_ws.append(
w([self.fc_layer_sizes[i],
self.fc_layer_sizes[i + 1]]))
fc_bs.append(b([self.fc_layer_sizes[i + 1]]))
##
# Forward pass calculation
##
def generate_predictions():
"""
Runs self.input_frames through the network to generate a prediction from 0
(generated img) to 1 (real img).
@return: A tensor of predictions of shape [self.batch_size x 1].
"""
with tf.name_scope('calculation'):
preds = tf.zeros([self.batch_size, 1])
last_input = self.input_frames
# convolutions
with tf.name_scope('convolutions'):
for i in xrange(len(conv_ws)):
# Convolve layer and activate with ReLU
preds = tf.nn.conv2d(last_input,
conv_ws[i], [1, 1, 1, 1],
padding=c.PADDING_D)
preds = tf.nn.relu(preds + conv_bs[i])
last_input = preds
# pooling layer
with tf.name_scope('pooling'):
preds = tf.nn.max_pool(preds, [1, 2, 2, 1], [1, 2, 2, 1],
padding=c.PADDING_D)
# flatten preds for dense layers
shape = preds.get_shape().as_list()
# -1 can be used as one dimension to size dynamically
preds = tf.reshape(
preds,
[-1, (c.HIST_LEN + 1), shape[1] * shape[2] * shape[3]])
print 'preds:', preds.get_shape()
# conv_outputs = tf.stack(tf.split(preds, self.batch_size), 1)
conv_outputs = tf.transpose(preds, [1, 0, 2])
print 'conv_outputs:', conv_outputs.get_shape()
print 'hidden_dim:', hidden_dim
# lstm layers
with tf.name_scope('lstm'):
lstm_cell = rnn.BasicLSTMCell(hidden_dim)
dropout = rnn.DropoutWrapper(lstm_cell,
output_keep_prob=c.KEEP_PROB)
stacked_lstm = rnn.MultiRNNCell([dropout] * c.LAYERS)
initial_state = stacked_lstm.zero_state(
self.batch_size, tf.float32)
outputs, state = tf.nn.dynamic_rnn(
stacked_lstm,
conv_outputs,
initial_state=initial_state,
time_major=True,
dtype=tf.float32,
scope='lstm_' + str(self.scale_index))
preds = outputs[-1]
# fully-connected layers
with tf.name_scope('fully-connected'):
for i in xrange(len(fc_ws)):
preds = tf.matmul(preds, fc_ws[i]) + fc_bs[i]
# Activate with ReLU (or Sigmoid for last layer)
if i == len(fc_ws) - 1:
preds = tf.sigmoid(preds)
else:
preds = tf.nn.relu(preds)
# clip preds between [.1, 0.9] for stability
with tf.name_scope('clip'):
preds = tf.clip_by_value(preds, 0.1, 0.9)
return preds
self.preds = generate_predictions()
def define_graph(self):
"""
Sets up the model graph in TensorFlow.
"""
self.train_vars = [] # the variables to train in the optimization step
##
# Layer setup
##
with tf.name_scope('setup'):
# convolution
with tf.name_scope('convolutions'):
self.conv_ws = []
self.conv_bs = []
last_out_height = self.height
last_out_width = self.width
with tf.name_scope('weights'):
for i in xrange(len(self.kernel_sizes)):
self.conv_ws.append(
w([
self.kernel_sizes[i], self.kernel_sizes[i],
self.conv_layer_fms[i],
self.conv_layer_fms[i + 1]
], 'dis_con_' + str(self.scale_index) + '_' +
str(i)))
with tf.name_scope('biases'):
for i in xrange(len(self.kernel_sizes)):
self.conv_bs.append(b([self.conv_layer_fms[i + 1]]))
last_out_height = conv_out_size(
last_out_height, self.padding,
self.kernel_sizes[i], 1)
last_out_width = conv_out_size(last_out_width,
self.padding,
self.kernel_sizes[i], 1)
# fully-connected
with tf.name_scope('fully-connected'):
# Add in an initial layer to go from the last conv to the first fully-connected.
# Use /2 for the height and width if not using the new models because there is a 2x2 pooling layer
self.fc_layer_sizes.insert(
0,
last_out_height * last_out_width * self.conv_layer_fms[-1])
self.fc_ws = []
self.fc_bs = []
with tf.name_scope('weights'):
for i in xrange(len(self.fc_layer_sizes) - 1):
self.fc_ws.append(
w([
self.fc_layer_sizes[i],
self.fc_layer_sizes[i + 1]
], 'dis_fc_' + str(self.scale_index) + '_' +
str(i)))
with tf.name_scope('biases'):
for i in xrange(len(self.fc_layer_sizes) - 1):
self.fc_bs.append(b([self.fc_layer_sizes[i + 1]]))
self.train_vars += self.conv_ws
self.train_vars += self.conv_bs
self.train_vars += self.fc_ws
self.train_vars += self.fc_bs
def define_graph(self):
"""
Define model graph.
"""
if self.verbose: print("Defining graph...")
##
# DEFINE INPUT DATA
##
if self.inverse:
self.x = self.input_features["gather"]
self.y_true = self.input_features["reflectivity"]
else:
self.x = self.input_features["reflectivity"]
self.y_true = self.input_features["gather"]
self.velocity = self.input_features["velocity"]
# INPUT/OUTPUT HAS SHAPE NWC
self.x_shape = self.x.shape.as_list()
self.y_true_shape = self.y_true.shape.as_list()
##
# DEFINE VARIABLES
##
# define weights for wavenet
self.W = Wavenet1D(in_channels=self.x_shape[2],
filter_width=2,
num_blocks=self.c.NUM_WAVE_BLOCKS,
num_layers=len(self.c.WAVE_RATES),
hidden_channels=self.c.WAVE_HIDDEN_CHANNELS,
rates=self.c.WAVE_RATES,
activation=self.c.WAVE_ACTIVATION,
biases=self.c.WAVE_BIASES,
verbose=False)
self.W.define_variables()
# define weights for final convolutional layer
self.CONV_KERNEL = [self.c.CONV_FILTER_LENGTH,self.c.WAVE_HIDDEN_CHANNELS,self.y_true_shape[2]]
self.weights, self.biases = {}, {}
with tf.name_scope('conv1d_params'):
stddev = np.sqrt(1) / np.sqrt(np.prod(self.CONV_KERNEL[:2]))
weights = w(self.CONV_KERNEL, mean=0., stddev=stddev, name="weights")
biases = b(self.CONV_KERNEL[2:], const=0.0, name="biases")
self.weights["conv1d"] = weights
self.biases["conv1d"] = biases
##
# DEFINE GRAPH
##
def construct_layers(x):
if self.verbose:
print("y_true: ",self.y_true.shape)
print("x: ",x.shape)
if self.inverse: x = x[:,::-1,:]# FLIP DATA TO REMAIN CAUSAL
# WAVENET
x = self.W.define_graph(x)
if self.verbose: print("wavenet: ",x.shape)
# CONVOLUTION
with tf.name_scope("conv1d"):
# causal convolution
with tf.name_scope("pad_left"):
x = tf.pad(x, [[0, 0], [(self.CONV_KERNEL[0]-1), 0], [0, 0]])# pad appropriate zeros on input
x = tf.nn.convolution(x, filter=self.weights["conv1d"], strides=[1], padding="VALID", data_format="NWC")
x = x + self.biases["conv1d"]
if self.verbose: print("conv1d: ",x.shape)
if self.inverse: x = x[:,::-1,:]# FLIP DATA TO REMAIN CAUSAL
return x
## initialise network
self.y = construct_layers(self.x)
assert self.y.shape.as_list() == self.y_true.shape.as_list()
# print out number of weights
self.num_weights = np.sum([self.weights[tensor].shape.num_elements() for tensor in self.weights])
self.num_biases = np.sum([self.biases[tensor].shape.num_elements() for tensor in self.biases])
self.total_num_trainable_params = self.num_weights+self.num_biases+self.W.num_weights+self.W.num_biases
if self.verbose: print(self)
# check no more trainable variables introduced
assert self.total_num_trainable_params == np.sum([tensor.shape.num_elements() for tensor in tf.trainable_variables()])
def define_graph(self, discriminator):
# Sets up the model graph
with tf.name_scope('generator'):
# Data
with tf.name_scope('data'):
# Prepare the placeholder for train input frames
self.input_frames_train = tf.placeholder(tf.float32, shape=[None, self.frame_height,
self.frame_width, 3*c.HIST_LEN])
# Prepare the placeholder for train ground-truth frames
self.gt_frames_train = tf.placeholder(tf.float32, shape=[None, self.frame_height,
self.frame_width, 3*c.OUT_LEN])
# Prepare the placeholder for test input frames
self.input_frames_test = tf.placeholder(tf.float32, shape=[None, self.frame_height,
self.frame_width, 3*c.HIST_LEN])
# Prepare the placeholder for test ground-truth frames
self.gt_frames_test = tf.placeholder(tf.float32, shape=[None, self.frame_height,
self.frame_width, 3*c.OUT_LEN])
# Variable batchsize
self.batch_size_train = tf.shape(self.input_frames_train)[0]
self.batch_size_test = tf.shape(self.input_frames_test)[0]
self.train_vars = [] # the variables to train
self.summaries_train = []
self.preds_train = [] # the generated images
self.gts_train = [] # the ground truth images
self.d_preds = [] # the predictions from the discriminator model
self.summaries_test = []
self.preds_test = [] # the generated images
self.gts_test = [] # the ground truth images
self.ws = []
self.bs = []
# Sets up the generator network
with tf.name_scope('scale_1'):
with tf.name_scope('setup'):
# Sets up the weights and biases
scale_ws = []
scale_bs = []
for i in xrange(len(self.kernel_sizes)):
scale_ws.append(w([self.kernel_sizes[i], self.kernel_sizes[i],
self.feature_maps[i],
self.feature_maps[i + 1]]))
scale_bs.append(b([self.feature_maps[i + 1]]))
# Add to trainable parameters
self.train_vars += scale_ws
self.train_vars += scale_bs
self.ws.append(scale_ws)
self.bs.append(scale_bs)
with tf.name_scope('calculation'):
# Perform train calculation
train_preds, train_gts = self.generate_predictions(self.input_frames_train, self.gt_frames_train)
self.preds_train.append(train_preds)
self.gts_train.append(train_gts)
test_preds, test_gts = self.generate_predictions(self.input_frames_test, self.gt_frames_test)
self.preds_test.append(test_preds)
self.gts_test.append(test_gts)
with tf.name_scope('d_preds'):
self.d_preds = []
with tf.name_scope('scale_1'):
with tf.name_scope('calculation'):
# IMPLEMENT THIS FIRST (INCOMPLETE)
self.d_preds.append(discriminator.nets.generate_predictions())
with tf.name_scope('train'):
self.global_loss = combined_loss(self.d_preds, self.input_frames_train, self.preds_train, self.gts_train)
self.global_step = tf.Variable(0, trainable=False)
self.optimizer = tf.train.AdamOptimizer(learning_rate=c.LRATE_G, name='optim')
self.train_op = self.optimizer.minimize(self.global_loss, global_step=self.global_step,
var_list=self.train_vars, name='train_op')
loss_summary = tf.summary.scalar('train_loss_G', self.global_loss)
self.summaries_train.append(loss_summary)
def define_graph(self):
"""
Sets up the model graph in TensorFlow.
"""
with tf.name_scope('generator'):
##
# Data
##
with tf.name_scope('data'):
self.inputs = tf.placeholder(tf.float32, shape=[None, 6])
self.gt_frames = tf.placeholder(
tf.float32, shape=[None, self.height, self.width, 3])
# use variable batch_size for more flexibility
self.batch_size = tf.shape(self.inputs)[0]
##
# Scale network setup and calculation
##
self.summaries = []
self.scale_preds = [] # the generated images at each scale
self.scale_gts = [] # the ground truth images at each scale
self.d_scale_preds = [
] # the predictions from the discriminator model
for scale_num in xrange(self.num_scale_nets):
with tf.name_scope('scale_' + str(scale_num)):
with tf.name_scope('setup'):
with tf.name_scope('fully-connected'):
fc_ws = []
fc_bs = []
# create weights for fc layers
for i in xrange(
len(self.scale_fc_layer_sizes[scale_num]) -
1):
fc_ws.append(
w([
self.scale_fc_layer_sizes[scale_num]
[i],
self.scale_fc_layer_sizes[scale_num][i
+
1]
]))
fc_bs.append(
b([
self.scale_fc_layer_sizes[scale_num][i
+
1]
]))
with tf.name_scope('convolutions'):
conv_ws = []
conv_bs = []
# create weights for kernels
for i in xrange(
len(self.scale_kernel_sizes[scale_num])):
conv_ws.append(
w([
self.scale_kernel_sizes[scale_num][i],
self.scale_kernel_sizes[scale_num][i],
self.scale_conv_layer_fms[scale_num]
[i],
self.scale_conv_layer_fms[scale_num][i
+
1]
]))
conv_bs.append(
b([
self.scale_conv_layer_fms[scale_num][i
+
1]
]))
with tf.name_scope('calculation'):
def calculate(height, width, inputs, gts,
last_gen_frames):
# scale inputs and gts
scale_factor = 1. / 2**(
(self.num_scale_nets - 1) - scale_num)
scale_height = int(height * scale_factor)
scale_width = int(width * scale_factor)
scale_gts = tf.image.resize_images(
gts, scale_height, scale_width)
# for all scales but the first, add the frame generated by the last
# scale to the input
# if scale_num > 0:
# last_gen_frames = tf.image.resize_images(last_gen_frames,
# scale_height,
# scale_width)
# inputs = tf.concat(3, [inputs, last_gen_frames])
# generated frame predictions
preds = inputs
# perform fc multiplications
with tf.name_scope('fully-connected'):
for i in xrange(
len(self.
scale_fc_layer_sizes[scale_num]) -
1):
preds = tf.nn.relu(
tf.matmul(preds, fc_ws[i]) + fc_bs[i])
# reshape for convolutions
preds = tf.reshape(preds, [
-1, c.FRAME_HEIGHT, c.FRAME_WIDTH,
self.scale_conv_layer_fms[scale_num][0]
])
# perform convolutions
with tf.name_scope('convolutions'):
for i in xrange(
len(self.scale_kernel_sizes[scale_num])
):
# Convolve layer
preds = tf.nn.conv2d(preds,
conv_ws[i],
[1, 1, 1, 1],
padding=c.PADDING_G)
# Activate with ReLU (or Tanh for last layer)
if i == len(
self.scale_kernel_sizes[scale_num]
) - 1:
preds = tf.nn.tanh(preds + conv_bs[i])
else:
preds = tf.nn.relu(preds + conv_bs[i])
return preds, scale_gts
##
# Perform train calculation
##
# for all scales but the first, add the frame generated by the last
# scale to the input
if scale_num > 0:
last_scale_pred = self.scale_preds[scale_num - 1]
else:
last_scale_pred = None
# calculate
train_preds, train_gts = calculate(
self.height, self.width, self.inputs,
self.gt_frames, last_scale_pred)
self.scale_preds.append(train_preds)
self.scale_gts.append(train_gts)
# We need to run the network first to get generated frames, run the
# discriminator on those frames to get d_scale_preds, then run this
# again for the loss optimization.
if c.ADVERSARIAL:
self.d_scale_preds.append(
tf.placeholder(tf.float32, [None, 1]))
##
# Training
##
with tf.name_scope('train'):
# global loss is the combined loss from every scale network
self.global_loss = combined_loss(self.scale_preds,
self.scale_gts,
self.d_scale_preds)
self.global_step = tf.Variable(0, trainable=False)
self.optimizer = tf.train.AdamOptimizer(
learning_rate=c.LRATE_G, name='optimizer')
self.train_op = self.optimizer.minimize(
self.global_loss,
global_step=self.global_step,
name='train_op')
# train loss summary
loss_summary = tf.scalar_summary('train_loss_G',
self.global_loss)
self.summaries.append(loss_summary)
##
# Error
##
with tf.name_scope('error'):
# error computation
# get error at largest scale
self.psnr_error = psnr_error(self.scale_preds[-1],
self.gt_frames)
self.sharpdiff_error = sharp_diff_error(
self.scale_preds[-1], self.gt_frames)
# train error summaries
summary_psnr = tf.scalar_summary('train_PSNR', self.psnr_error)
summary_sharpdiff = tf.scalar_summary('train_SharpDiff',
self.sharpdiff_error)
self.summaries += [summary_psnr, summary_sharpdiff]
# add summaries to visualize in TensorBoard
self.summaries = tf.merge_summary(self.summaries)
def define_graph(self, discriminator):
"""
Sets up the model graph in TensorFlow.
@param discriminator: The discriminator model that discriminates frames generated by this
model.
"""
with tf.name_scope('generator'):
##
# Data
##
with tf.name_scope('input'):
self.input_frames_train = tf.placeholder(
tf.float32,
shape=[
None, self.height_train, self.width_train,
3 * c.HIST_LEN
],
name='input_frames_train')
self.gt_frames_train = tf.placeholder(tf.float32,
shape=[
None,
self.height_train,
self.width_train,
3 * c.GT_LEN
],
name='gt_frames_train')
self.input_frames_test = tf.placeholder(
tf.float32,
shape=[
None, self.height_test, self.width_test, 3 * c.HIST_LEN
],
name='input_frames_test')
self.gt_frames_test = tf.placeholder(tf.float32,
shape=[
None,
self.height_test,
self.width_test,
3 * c.GT_LEN
],
name='gt_frames_test')
# use variable batch_size for more flexibility
with tf.name_scope('batch_size_train'):
self.batch_size_train = tf.shape(
self.input_frames_train,
name='input_frames_train_shape')[0]
with tf.name_scope('batch_size_test'):
self.batch_size_test = tf.shape(
self.input_frames_test,
name='input_frames_test_shape')[0]
##
# Scale network setup and calculation
##
self.train_vars = [
] # the variables to train in the optimization step
self.summaries_train = []
self.scale_preds_train = [] # the generated images at each scale
self.scale_gts_train = [] # the ground truth images at each scale
self.d_scale_preds = [
] # the predictions from the discriminator model
self.summaries_test = []
self.scale_preds_test = [] # the generated images at each scale
self.scale_gts_test = [] # the ground truth images at each scale
self.ws = []
self.bs = []
for scale_num in xrange(self.num_scale_nets):
with tf.name_scope('scale_net_' + str(scale_num)):
with tf.name_scope('setup'):
scale_ws = []
scale_bs = []
# create weights for kernels
with tf.name_scope('weights'):
for i in xrange(
len(self.scale_kernel_sizes[scale_num])):
scale_ws.append(
w([
self.scale_kernel_sizes[scale_num][i],
self.scale_kernel_sizes[scale_num][i],
self.scale_layer_fms[scale_num][i],
self.scale_layer_fms[scale_num][i + 1]
], 'gen_' + str(scale_num) + '_' + str(i)))
with tf.name_scope('biases'):
for i in xrange(
len(self.scale_kernel_sizes[scale_num])):
scale_bs.append(
b([self.scale_layer_fms[scale_num][i + 1]
]))
# add to trainable parameters
self.train_vars += scale_ws
self.train_vars += scale_bs
self.ws.append(scale_ws)
self.bs.append(scale_bs)
with tf.name_scope('calculation'):
with tf.name_scope('calculation_train'):
##
# Perform train calculation
##
if scale_num > 0:
last_scale_pred_train = self.scale_preds_train[
scale_num - 1]
else:
last_scale_pred_train = None
train_preds, train_gts = self.generate_predictions(
scale_num, self.height_train, self.width_train,
self.input_frames_train, self.gt_frames_train,
last_scale_pred_train)
with tf.name_scope('calculation_test'):
##
# Perform test calculation
if scale_num > 0:
last_scale_pred_test = self.scale_preds_test[
scale_num - 1]
else:
last_scale_pred_test = None
test_preds, test_gts = self.generate_predictions(
scale_num, self.height_test, self.width_test,
self.input_frames_test, self.gt_frames_test,
last_scale_pred_test, 'test')
self.scale_preds_train.append(train_preds)
self.scale_gts_train.append(train_gts)
self.scale_preds_test.append(test_preds)
self.scale_gts_test.append(test_gts)
##
# Get Discriminator Predictions
##
if c.ADVERSARIAL:
with tf.name_scope('d_preds'):
# A list of the prediction tensors for each scale network
self.d_scale_preds = []
for scale_num in xrange(self.num_scale_nets):
with tf.name_scope('scale_' + str(scale_num)):
with tf.name_scope('calculation'):
input_scale_factor = 1. / self.scale_gt_inverse_scale_factor[
scale_num]
input_scale_height = int(self.height_train *
input_scale_factor)
input_scale_width = int(self.width_train *
input_scale_factor)
scale_inputs_train = tf.image.resize_images(
self.input_frames_train,
[input_scale_height, input_scale_width])
# get predictions from the d scale networks
self.d_scale_preds.append(
discriminator.scale_nets[scale_num].
generate_all_predictions(
scale_inputs_train,
self.scale_preds_train[scale_num]))
##
# Training
##
with tf.name_scope('training'):
# global loss is the combined loss from every scale network
self.global_loss = temporal_combined_loss(
self.scale_preds_train, self.scale_gts_train,
self.d_scale_preds)
with tf.name_scope('train_step'):
self.global_step = tf.Variable(0,
trainable=False,
name='global_step')
self.optimizer = tf.train.AdamOptimizer(
learning_rate=c.LRATE_G, name='optimizer')
self.train_op = self.optimizer.minimize(
self.global_loss,
global_step=self.global_step,
var_list=self.train_vars,
name='train_op')
# train loss summary
loss_summary = tf.summary.scalar('train_loss_G',
self.global_loss)
self.summaries_train.append(loss_summary)
##
# Error
##
with tf.name_scope('error'):
# error computation
# get error at largest scale
with tf.name_scope('psnr_train'):
self.psnr_error_train = []
for gt_num in xrange(c.GT_LEN):
self.psnr_error_train.append(
psnr_error(
self.scale_preds_train[-1][:, :, :, gt_num *
3:(gt_num + 1) * 3],
self.gt_frames_train[:, :, :, gt_num *
3:(gt_num + 1) * 3]))
with tf.name_scope('sharpdiff_train'):
self.sharpdiff_error_train = []
for gt_num in xrange(c.GT_LEN):
self.sharpdiff_error_train.append(
sharp_diff_error(
self.scale_preds_train[-1][:, :, :, gt_num *
3:(gt_num + 1) * 3],
self.gt_frames_train[:, :, :, gt_num *
3:(gt_num + 1) * 3]))
with tf.name_scope('ssim_train'):
self.ssim_error_train = []
for gt_num in xrange(c.GT_LEN):
self.ssim_error_train.append(
ssim_error(
self.scale_preds_train[-1][:, :, :, gt_num *
3:(gt_num + 1) * 3],
self.gt_frames_train[:, :, :, gt_num *
3:(gt_num + 1) * 3]))
with tf.name_scope('psnr_test'):
self.psnr_error_test = []
for gt_num in xrange(c.GT_LEN):
self.psnr_error_test.append(
psnr_error(
self.scale_preds_test[-1][:, :, :, gt_num *
3:(gt_num + 1) * 3],
self.gt_frames_test[:, :, :, gt_num *
3:(gt_num + 1) * 3]))
with tf.name_scope('sharpdiff_test'):
self.sharpdiff_error_test = []
for gt_num in xrange(c.GT_LEN):
self.sharpdiff_error_test.append(
sharp_diff_error(
self.scale_preds_test[-1][:, :, :, gt_num *
3:(gt_num + 1) * 3],
self.gt_frames_test[:, :, :, gt_num *
3:(gt_num + 1) * 3]))
with tf.name_scope('ssim_test'):
self.ssim_error_test = []
for gt_num in xrange(c.GT_LEN):
self.ssim_error_test.append(
ssim_error(
self.scale_preds_test[-1][:, :, :, gt_num *
3:(gt_num + 1) * 3],
self.gt_frames_test[:, :, :, gt_num *
3:(gt_num + 1) * 3]))
for gt_num in xrange(c.GT_LEN):
# train error summaries
summary_psnr_train = tf.summary.scalar(
'train_PSNR_' + str(gt_num),
self.psnr_error_train[gt_num])
summary_sharpdiff_train = tf.summary.scalar(
'train_SharpDiff_' + str(gt_num),
self.sharpdiff_error_train[gt_num])
summary_ssim_train = tf.summary.scalar(
'train_SSIM_' + str(gt_num),
self.ssim_error_train[gt_num])
self.summaries_train += [
summary_psnr_train, summary_sharpdiff_train,
summary_ssim_train
]
# test error summaries
summary_psnr_test = tf.summary.scalar(
'test_PSNR_' + str(gt_num),
self.psnr_error_test[gt_num])
summary_sharpdiff_test = tf.summary.scalar(
'test_SharpDiff_' + str(gt_num),
self.sharpdiff_error_test[gt_num])
summary_ssim_test = tf.summary.scalar(
'test_SSIM_' + str(gt_num),
self.ssim_error_test[gt_num])
self.summaries_test += [
summary_psnr_test, summary_sharpdiff_test,
summary_ssim_test
]
# add summaries to visualize in TensorBoard
self.summaries_train = tf.summary.merge(self.summaries_train)
self.summaries_test = tf.summary.merge(self.summaries_test)
