snapshot...

modules/deeplearning/espcn.py

@@ -361,9 +361,8 @@ class ESPCN:
 
         self.get_evaluate_dataset(idxs)
 
-    def build_espcn(self, do_drop_out=False, do_batch_norm=False, drop_rate=0.5):
+    def build_espcn(self, do_drop_out=False, do_batch_norm=False, drop_rate=0.5, factor=2):
         print('build_cnn')
-        # padding = "VALID"
         padding = "SAME"
 
         # activation = tf.nn.relu
@@ -371,78 +370,36 @@ class ESPCN:
         activation = tf.nn.leaky_relu
         momentum = 0.99
 
-        # num_filters = len(self.train_params) * 4
-        num_filters = self.n_chans * 64
+        num_filters = 64
 
         input_2d = self.inputs[0]
         print('input: ', input_2d.shape)
+
         # conv = tf.keras.layers.Conv2D(num_filters, kernel_size=5, strides=1, padding='VALID', activation=None)(input_2d)
         conv = input_2d
         print('input: ', conv.shape)
         skip = conv
 
+        conv = conv_b = tf.keras.layers.Conv2D(num_filters, kernel_size=3, padding=padding)(input_2d)
+
         if NOISE_TRAINING:
             conv = tf.keras.layers.GaussianNoise(stddev=NOISE_STDDEV)(conv)
 
-        if do_drop_out:
-            conv = tf.keras.layers.Dropout(drop_rate)(conv)
-        if do_batch_norm:
-            conv = tf.keras.layers.BatchNormalization()(conv)
-
-        conv = tf.keras.layers.Conv2D(num_filters, kernel_size=5, strides=1, padding=padding, activation=activation)(conv)
-        print(conv.shape)
-
-        if do_drop_out:
-            conv = tf.keras.layers.Dropout(drop_rate)(conv)
-        if do_batch_norm:
-            conv = tf.keras.layers.BatchNormalization()(conv)
-
-        conv = tf.keras.layers.Conv2D(num_filters, kernel_size=3, strides=1, padding=padding, activation=activation)(conv)
-        print(conv.shape)
-
-        # conv = tf.keras.layers.Conv2D(num_filters, kernel_size=3, strides=1, padding=padding, activation=None)(conv)
-        # conv = tf.keras.layers.BatchNormalization()(conv)
-        # print(conv.shape)
-        #
-        # conv = conv + skip
-        # conv = tf.keras.layers.LeakyReLU()(conv)
-        # print(conv.shape)
-
-        if do_drop_out:
-            conv = tf.keras.layers.Dropout(drop_rate)(conv)
-        if do_batch_norm:
-            conv = tf.keras.layers.BatchNormalization()(conv)
-
-        conv = tf.keras.layers.Conv2D(num_filters // 2, kernel_size=3, strides=1, padding=padding, activation=activation)(conv)
-        conv = tf.keras.layers.BatchNormalization()(conv)
-        print(conv.shape)
+        conv_b = build_conv2d_block(conv_b, num_filters, 'Residual_Block_1')
 
-        if do_drop_out:
-            conv = tf.keras.layers.Dropout(drop_rate)(conv)
-        if do_batch_norm:
-            conv = tf.keras.layers.BatchNormalization()(conv)
+        conv_b = build_conv2d_block(conv_b, num_filters, 'Residual_Block_2')
 
-        # conv = tf.keras.layers.Conv2D(num_filters // 2, kernel_size=3, strides=1, padding=padding, activation=activation)(conv)
-        # print(conv.shape)
+        conv_b = build_conv2d_block(conv_b, num_filters, 'Residual_Block_3')
 
-        conv = tf.keras.layers.Conv2DTranspose(num_filters // 4, kernel_size=4, strides=2, padding=padding, activation=activation)(conv)
-        print(conv.shape)
+        conv_b = tf.keras.layers.Conv2D(num_filters, kernel_size=3, strides=1, padding=padding)(conv_b)
 
-        conv = tf.keras.layers.Conv2DTranspose(num_filters // 4, kernel_size=1, strides=1, padding=padding, activation=activation)(conv)
-        print(conv.shape)
+        conv = conv + conv_b
 
-        conv = tf.keras.layers.Conv2DTranspose(num_filters // 4, kernel_size=1, strides=1, padding=padding, activation=activation)(conv)
-        print(conv.shape)
+        conv = tf.keras.layers.Conv2D(num_filters * (factor ** 2), 3, padding='same')(conv)
 
-        #self.logits = tf.keras.layers.Conv2D(1, kernel_size=1, strides=1, padding=padding, name='probability', activation=tf.nn.sigmoid)(conv)
-        self.logits = tf.keras.layers.Conv2D(1, kernel_size=1, strides=1, padding=padding, name='probability')(conv)
+        conv = tf.nn.depth_to_space(conv, factor)
 
-        # conv = tf.nn.depth_to_space(conv, block_size=2)
-        # conv = tf.keras.layers.Activation(activation=activation)(conv)
-        # print(conv.shape)
-        #
-        # # Called logits, but these are actually probabilities, see activation
-        # self.logits = tf.keras.layers.Activation(activation=activation)(conv)
+        self.logits = tf.keras.layers.Conv2D(1, kernel_size=3, strides=1, padding=padding, name='regression')(conv)
 
         print(self.logits.shape)