diff --git a/video_prediction_tools/model_modules/video_prediction/layers/layer_def.py b/video_prediction_tools/model_modules/video_prediction/layers/layer_def.py
index 925f94c44156f86fe43f168e63b3bb5935c2db04..79d7653a5cc55d72abb7ea2fbdb22ca8be4c3b67 100644
--- a/video_prediction_tools/model_modules/video_prediction/layers/layer_def.py
+++ b/video_prediction_tools/model_modules/video_prediction/layers/layer_def.py
@@ -6,7 +6,6 @@
"""
import tensorflow as tf
-import numpy as np
weight_decay = 0.0005
def _activation_summary(x):
@@ -67,7 +66,7 @@ def conv_layer(inputs, kernel_size, stride, num_features, idx, initializer=tf.co
input_channels, num_features],
stddev = 0.01, wd = weight_decay)
biases = _variable_on_gpu('biases', [num_features], initializer)
- conv = tf.nn.conv2d(inputs, weights, strides = [1, stride, stride, 1], padding = 'SAME')
+ conv = tf.nn.conv2d(inputs, weights, strides = [1, stride, stride, 1], padding='SAME')
conv_biased = tf.nn.bias_add(conv, biases)
if activate == "linear":
return conv_biased
@@ -162,4 +161,3 @@ def bn_layers(inputs,idx,is_training=True,epsilon=1e-3,decay=0.99,reuse=None):
def bn_layers_wrapper(inputs, is_training):
pass
-
diff --git a/video_prediction_tools/model_modules/video_prediction/models/weatherBench3DCNN.py b/video_prediction_tools/model_modules/video_prediction/models/weatherBench3DCNN.py
index ac172074484648871f586b9aeed868ea6966b3fa..07b5cbf26251e5d61520b07cc25e22d64338d195 100644
--- a/video_prediction_tools/model_modules/video_prediction/models/weatherBench3DCNN.py
+++ b/video_prediction_tools/model_modules/video_prediction/models/weatherBench3DCNN.py
@@ -77,9 +77,9 @@ class WeatherBenchModel(object):
original_global_variables = tf.global_variables()
# Architecture
- x_hat = self.build_model(x, self.filters, self.kernels, dr=0)
+ x_hat = self.build_model(x, self.filters, self.kernels)
# Loss
- self.total_loss = l1_loss(x[...,0], x_hat[...,0])
+ self.total_loss = l1_loss(x[, 0,:, :,0], x_hat[, 0,:, :,0])
# Optimizer
self.train_op = tf.train.AdamOptimizer(
@@ -97,47 +97,14 @@ class WeatherBenchModel(object):
return self.is_build_graph
- def build_model(self, x, filters, kernels, dr=0):
+ def build_model(self, x, filters, kernels):
"""Fully convolutional network"""
for f, k in zip(filters[:-1], kernels[:-1]):
- x = PeriodicConv2D(x, f, k)
- x = tf.nn.elu(x)
- if dr > 0: x = tf.nn.dropout(x, dr)
- output = PeriodicConv2D(x, filters[-1], kernels[-1])
+ x = ld.conv_layer(x, kernel_size=k, stride=1, num_features=f, idx="Conv_layer", activate="leaky_relu")
+ output = ld.conv_layer(x, kernel_size=kernels[-1], stride=1, num_features=filters[-1], idx="Conv_layer", activate="linear")
return output
-class PeriodicPadding2D(object):
- def __init__(self, pad_width):
-
- self.pad_width = pad_width
-
- def call(self, inputs, **kwargs):
- if self.pad_width == 0:
- return inputs
-
- inputs_padded = tf.concat(
- [inputs[:, :, -self.pad_width:, :], inputs, inputs[:, :, :self.pad_width, :]], axis=2)
-
- # Zero padding in the lat direction
- inputs_padded = tf.pad(inputs_padded, [[0, 0], [self.pad_width, self.pad_width], [0, 0], [0, 0]])
- return inputs_padded
-
-
-class PeriodicConv2D(object):
-
- def __init__(self, filters, kernel_size, conv_kwargs={}):
- self.filters = filters
- self.kernel_size = kernel_size
- self.conv_kwargs = conv_kwargs
- if type(kernel_size) is not int:
- assert kernel_size[0] == kernel_size[1], 'PeriodicConv2D only works for square kernels'
- kernel_size = kernel_size[0]
- self.pad_width = (kernel_size - 1) // 2
-
- def call(self,inputs):
- self.padding = PeriodicPadding2D(inputs, self.pad_width)
- self.conv = ld.conv2D(self.padding, self.filters, self.kernel_size, padding='valid')