def __init__(self, input_dim, output_dim=1, optimizer=None, use_dropout=True, lr=0.001, random_seed=123,
                 is_train=True, log_dir=None, name=None):
        self.name = name
        self.is_train = is_train
        self.log_dir = log_dir
        self.cur_lr = None
        self.logger, self.file_handler, self.stream_handler = utils.init_logger(log_dir=self.log_dir,
                                                                                name=self.name,
                                                                                is_train=self.is_train)
        with tf.variable_scope(self.name):
            # Placeholders for inputs
            self.X = tf.placeholder(dtype=tf.float32, shape=[None, input_dim], name='X')
            self.y = tf.placeholder(dtype=tf.float32, shape=[None, output_dim], name='y')
            self.keep_prob = tf.placeholder(tf.float32, name='keep_prob')
            tf_utils.print_activations(self.X, logger=self.logger if self.is_train else None)

            # Placeholders for TensorBoard
            self.train_acc = tf.placeholder(tf.float32, name='train_acc')
            self.val_acc = tf.placeholder(tf.float32, name='val_acc')

            net = self.X
            if use_dropout:
                net = tf_utils.dropout(x=net,
                                       keep_prob=self.keep_prob,
                                       seed=random_seed,
                                       name='dropout',
                                       logger=self.logger if self.is_train else None)

            # Network, loss, and optimizer
            self.y_pred = tf_utils.linear(net, output_size=output_dim)
            tf_utils.print_activations(self.y_pred, logger=self.logger if self.is_train else None)
            self.loss = tf.math.reduce_mean(tf.nn.l2_loss(self.y_pred - self.y))
            self.train_op, self.cur_lr = optimizer_fn(optimizer, lr=lr, loss=self.loss, name=self.name)

            # Accuracy etc
            self.y_pred_round = tf.math.round(x=self.y_pred, name='rounded_pred')
            accuracy = tf.equal(tf.cast(x=self.y_pred_round, dtype=tf.int32), tf.cast(x=self.y, dtype=tf.int32))
            self.accuracy = tf.reduce_mean(tf.cast(x=accuracy, dtype=tf.float32)) * 100.

        self._tensorboard()
        tf_utils.show_all_variables(logger=self.logger if self.is_train else None)
Пример #2
0
    def forward_network(self, img, padding='SAME', reuse=False):
        with tf.compat.v1.variable_scope(self.name, reuse=reuse):
            # Stage 0
            s0_conv1 = tf_utils.conv2d(x=img,
                                       output_dim=self.conv_dims[0],
                                       k_h=3,
                                       k_w=3,
                                       d_h=1,
                                       d_w=1,
                                       padding=padding,
                                       initializer='He',
                                       name='s0_conv1')
            s0_conv1 = tf_utils.relu(s0_conv1, name='relu_s0_conv1')

            s0_conv2 = tf_utils.conv2d(x=s0_conv1,
                                       output_dim=self.conv_dims[0],
                                       k_h=3,
                                       k_w=3,
                                       d_h=1,
                                       d_w=1,
                                       padding=padding,
                                       initializer='He',
                                       name='s0_conv2')
            s0_conv2 = tf_utils.norm(s0_conv2,
                                     name='s0_norm1',
                                     _type='batch',
                                     _ops=self._ops,
                                     is_train=False)
            s0_conv2 = tf_utils.relu(s0_conv2, name='relu_s0_conv2')

            # Stage 1
            s1_maxpool = tf_utils.max_pool(x=s0_conv2, name='s1_maxpool2d')

            s1_conv1 = tf_utils.conv2d(x=s1_maxpool,
                                       output_dim=self.conv_dims[0],
                                       k_h=3,
                                       k_w=3,
                                       d_h=1,
                                       d_w=1,
                                       padding=padding,
                                       initializer='He',
                                       name='s1_conv1')
            s1_conv1 = tf_utils.norm(s1_conv1,
                                     name='s1_norm0',
                                     _type='batch',
                                     _ops=self._ops,
                                     is_train=False)
            s1_conv1 = tf_utils.relu(s1_conv1, name='relu_s1_conv1')

            s1_conv2 = tf_utils.conv2d(x=s1_conv1,
                                       output_dim=self.conv_dims[1],
                                       k_h=3,
                                       k_w=3,
                                       d_h=1,
                                       d_w=1,
                                       padding=padding,
                                       initializer='He',
                                       name='s1_conv2')
            s1_conv2 = tf_utils.norm(s1_conv2,
                                     name='s1_norm1',
                                     _type='batch',
                                     _ops=self._ops,
                                     is_train=False)
            s1_conv2 = tf_utils.relu(s1_conv2, name='relu_s1_conv2')

            # Stage 2
            s2_maxpool = tf_utils.max_pool(x=s1_conv2, name='s2_maxpool2d')
            s2_conv1 = tf_utils.conv2d(x=s2_maxpool,
                                       output_dim=self.conv_dims[2],
                                       k_h=3,
                                       k_w=3,
                                       d_h=1,
                                       d_w=1,
                                       padding=padding,
                                       initializer='He',
                                       name='s2_conv1')
            s2_conv1 = tf_utils.norm(s2_conv1,
                                     name='s2_norm0',
                                     _type='batch',
                                     _ops=self._ops,
                                     is_train=False)
            s2_conv1 = tf_utils.relu(s2_conv1, name='relu_s2_conv1')

            s2_conv2 = tf_utils.conv2d(x=s2_conv1,
                                       output_dim=self.conv_dims[3],
                                       k_h=3,
                                       k_w=3,
                                       d_h=1,
                                       d_w=1,
                                       padding=padding,
                                       initializer='He',
                                       name='s2_conv2')
            s2_conv2 = tf_utils.norm(s2_conv2,
                                     name='s2_norm1',
                                     _type='batch',
                                     _ops=self._ops,
                                     is_train=False)
            s2_conv2 = tf_utils.relu(s2_conv2, name='relu_s2_conv2')

            # Stage 3
            s3_maxpool = tf_utils.max_pool(x=s2_conv2, name='s3_maxpool2d')
            s3_conv1 = tf_utils.conv2d(x=s3_maxpool,
                                       output_dim=self.conv_dims[4],
                                       k_h=3,
                                       k_w=3,
                                       d_h=1,
                                       d_w=1,
                                       padding=padding,
                                       initializer='He',
                                       name='s3_conv1')
            s3_conv1 = tf_utils.norm(s3_conv1,
                                     name='s3_norm0',
                                     _type='batch',
                                     _ops=self._ops,
                                     is_train=False)
            s3_conv1 = tf_utils.relu(s3_conv1, name='relu_s3_conv1')

            s3_conv2 = tf_utils.conv2d(x=s3_conv1,
                                       output_dim=self.conv_dims[5],
                                       k_h=3,
                                       k_w=3,
                                       d_h=1,
                                       d_w=1,
                                       padding=padding,
                                       initializer='He',
                                       name='s3_conv2')
            s3_conv2 = tf_utils.norm(s3_conv2,
                                     name='s3_norm1',
                                     _type='batch',
                                     _ops=self._ops,
                                     is_train=False)
            s3_conv2 = tf_utils.relu(s3_conv2, name='relu_s3_conv2')

            # Stage 4
            s4_maxpool = tf_utils.max_pool(x=s3_conv2, name='s4_maxpool2d')
            s4_conv1 = tf_utils.conv2d(x=s4_maxpool,
                                       output_dim=self.conv_dims[6],
                                       k_h=3,
                                       k_w=3,
                                       d_h=1,
                                       d_w=1,
                                       padding=padding,
                                       initializer='He',
                                       name='s4_conv1')
            s4_conv1 = tf_utils.norm(s4_conv1,
                                     name='s4_norm0',
                                     _type='batch',
                                     _ops=self._ops,
                                     is_train=False)
            s4_conv1 = tf_utils.relu(s4_conv1, name='relu_s4_conv1')

            s4_conv2 = tf_utils.conv2d(x=s4_conv1,
                                       output_dim=self.conv_dims[7],
                                       k_h=3,
                                       k_w=3,
                                       d_h=1,
                                       d_w=1,
                                       padding=padding,
                                       initializer='He',
                                       name='s4_conv2')
            s4_conv2 = tf_utils.norm(s4_conv2,
                                     name='s4_norm1',
                                     _type='batch',
                                     _ops=self._ops,
                                     is_train=False)
            s4_conv2 = tf_utils.relu(s4_conv2, name='relu_s4_conv2')
            s4_conv2_drop = tf_utils.dropout(x=s4_conv2,
                                             keep_prob=0.,
                                             name='s4_dropout')

            # Stage 5
            s5_maxpool = tf_utils.max_pool(x=s4_conv2_drop,
                                           name='s5_maxpool2d')
            s5_conv1 = tf_utils.conv2d(x=s5_maxpool,
                                       output_dim=self.conv_dims[8],
                                       k_h=3,
                                       k_w=3,
                                       d_h=1,
                                       d_w=1,
                                       padding=padding,
                                       initializer='He',
                                       name='s5_conv1')
            s5_conv1 = tf_utils.norm(s5_conv1,
                                     name='s5_norm0',
                                     _type='batch',
                                     _ops=self._ops,
                                     is_train=False)
            s5_conv1 = tf_utils.relu(s5_conv1, name='relu_s5_conv1')

            s5_conv2 = tf_utils.conv2d(x=s5_conv1,
                                       output_dim=self.conv_dims[9],
                                       k_h=3,
                                       k_w=3,
                                       d_h=1,
                                       d_w=1,
                                       padding=padding,
                                       initializer='He',
                                       name='s5_conv2')
            s5_conv2 = tf_utils.norm(s5_conv2,
                                     name='s5_norm1',
                                     _type='batch',
                                     _ops=self._ops,
                                     is_train=False)
            s5_conv2 = tf_utils.relu(s5_conv2, name='relu_s5_conv2')
            s5_conv2_drop = tf_utils.dropout(x=s5_conv2,
                                             keep_prob=0.,
                                             name='s5_dropout')

            # Stage 6
            s6_deconv1 = tf_utils.deconv2d(x=s5_conv2_drop,
                                           output_dim=self.conv_dims[10],
                                           k_h=2,
                                           k_w=2,
                                           initializer='He',
                                           name='s6_deconv1')
            s6_deconv1 = tf_utils.norm(s6_deconv1,
                                       name='s6_norm0',
                                       _type='batch',
                                       _ops=self._ops,
                                       is_train=False)
            s6_deconv1 = tf_utils.relu(s6_deconv1, name='relu_s6_deconv1')
            # Cropping
            w1 = s4_conv2_drop.get_shape().as_list()[2]
            w2 = s6_deconv1.get_shape().as_list()[2] - s4_conv2_drop.get_shape(
            ).as_list()[2]
            s6_deconv1_split, _ = tf.split(s6_deconv1,
                                           num_or_size_splits=[w1, w2],
                                           axis=2,
                                           name='axis2_split')
            tf_utils.print_activations(s6_deconv1_split)
            # Concat
            s6_concat = tf_utils.concat(
                values=[s6_deconv1_split, s4_conv2_drop],
                axis=3,
                name='s6_axis3_concat')

            s6_conv2 = tf_utils.conv2d(x=s6_concat,
                                       output_dim=self.conv_dims[11],
                                       k_h=3,
                                       k_w=3,
                                       d_h=1,
                                       d_w=1,
                                       padding=padding,
                                       initializer='He',
                                       name='s6_conv2')
            s6_conv2 = tf_utils.norm(s6_conv2,
                                     name='s6_norm1',
                                     _type='batch',
                                     _ops=self._ops,
                                     is_train=False)
            s6_conv2 = tf_utils.relu(s6_conv2, name='relu_s6_conv2')

            s6_conv3 = tf_utils.conv2d(x=s6_conv2,
                                       output_dim=self.conv_dims[12],
                                       k_h=3,
                                       k_w=3,
                                       d_h=1,
                                       d_w=1,
                                       padding=padding,
                                       initializer='He',
                                       name='s6_conv3')
            s6_conv3 = tf_utils.norm(s6_conv3,
                                     name='s6_norm2',
                                     _type='batch',
                                     _ops=self._ops,
                                     is_train=False)
            s6_conv3 = tf_utils.relu(s6_conv3, name='relu_s6_conv3')

            # Stage 7
            s7_deconv1 = tf_utils.deconv2d(x=s6_conv3,
                                           output_dim=self.conv_dims[13],
                                           k_h=2,
                                           k_w=2,
                                           initializer='He',
                                           name='s7_deconv1')
            s7_deconv1 = tf_utils.norm(s7_deconv1,
                                       name='s7_norm0',
                                       _type='batch',
                                       _ops=self._ops,
                                       is_train=False)
            s7_deconv1 = tf_utils.relu(s7_deconv1, name='relu_s7_deconv1')
            # Concat
            s7_concat = tf_utils.concat(values=[s7_deconv1, s3_conv2],
                                        axis=3,
                                        name='s7_axis3_concat')

            s7_conv2 = tf_utils.conv2d(x=s7_concat,
                                       output_dim=self.conv_dims[14],
                                       k_h=3,
                                       k_w=3,
                                       d_h=1,
                                       d_w=1,
                                       padding=padding,
                                       initializer='He',
                                       name='s7_conv2')
            s7_conv2 = tf_utils.norm(s7_conv2,
                                     name='s7_norm1',
                                     _type='batch',
                                     _ops=self._ops,
                                     is_train=False)
            s7_conv2 = tf_utils.relu(s7_conv2, name='relu_s7_conv2')

            s7_conv3 = tf_utils.conv2d(x=s7_conv2,
                                       output_dim=self.conv_dims[15],
                                       k_h=3,
                                       k_w=3,
                                       d_h=1,
                                       d_w=1,
                                       padding=padding,
                                       initializer='He',
                                       name='s7_conv3')
            s7_conv3 = tf_utils.norm(s7_conv3,
                                     name='s7_norm2',
                                     _type='batch',
                                     _ops=self._ops,
                                     is_train=False)
            s7_conv3 = tf_utils.relu(s7_conv3, name='relu_s7_conv3')

            # Stage 8
            s8_deconv1 = tf_utils.deconv2d(x=s7_conv3,
                                           output_dim=self.conv_dims[16],
                                           k_h=2,
                                           k_w=2,
                                           initializer='He',
                                           name='s8_deconv1')
            s8_deconv1 = tf_utils.norm(s8_deconv1,
                                       name='s8_norm0',
                                       _type='batch',
                                       _ops=self._ops,
                                       is_train=False)
            s8_deconv1 = tf_utils.relu(s8_deconv1, name='relu_s8_deconv1')
            # Concat
            s8_concat = tf_utils.concat(values=[s8_deconv1, s2_conv2],
                                        axis=3,
                                        name='s8_axis3_concat')

            s8_conv2 = tf_utils.conv2d(x=s8_concat,
                                       output_dim=self.conv_dims[17],
                                       k_h=3,
                                       k_w=3,
                                       d_h=1,
                                       d_w=1,
                                       padding=padding,
                                       initializer='He',
                                       name='s8_conv2')
            s8_conv2 = tf_utils.norm(s8_conv2,
                                     name='s8_norm1',
                                     _type='batch',
                                     _ops=self._ops,
                                     is_train=False)
            s8_conv2 = tf_utils.relu(s8_conv2, name='relu_s8_conv2')

            s8_conv3 = tf_utils.conv2d(x=s8_conv2,
                                       output_dim=self.conv_dims[18],
                                       k_h=3,
                                       k_w=3,
                                       d_h=1,
                                       d_w=1,
                                       padding=padding,
                                       initializer='He',
                                       name='s8_conv3')
            s8_conv3 = tf_utils.norm(s8_conv3,
                                     name='s8_norm2',
                                     _type='batch',
                                     _ops=self._ops,
                                     is_train=False)
            s8_conv3 = tf_utils.relu(s8_conv3, name='relu_conv3')

            # Stage 9
            s9_deconv1 = tf_utils.deconv2d(x=s8_conv3,
                                           output_dim=self.conv_dims[19],
                                           k_h=2,
                                           k_w=2,
                                           initializer='He',
                                           name='s9_deconv1')
            s9_deconv1 = tf_utils.norm(s9_deconv1,
                                       name='s9_norm0',
                                       _type='batch',
                                       _ops=self._ops,
                                       is_train=False)
            s9_deconv1 = tf_utils.relu(s9_deconv1, name='relu_s9_deconv1')
            # Concat
            s9_concat = tf_utils.concat(values=[s9_deconv1, s1_conv2],
                                        axis=3,
                                        name='s9_axis3_concat')

            s9_conv2 = tf_utils.conv2d(x=s9_concat,
                                       output_dim=self.conv_dims[20],
                                       k_h=3,
                                       k_w=3,
                                       d_h=1,
                                       d_w=1,
                                       padding=padding,
                                       initializer='He',
                                       name='s9_conv2')
            s9_conv2 = tf_utils.norm(s9_conv2,
                                     name='s9_norm1',
                                     _type='batch',
                                     _ops=self._ops,
                                     is_train=False)
            s9_conv2 = tf_utils.relu(s9_conv2, name='relu_s9_conv2')

            s9_conv3 = tf_utils.conv2d(x=s9_conv2,
                                       output_dim=self.conv_dims[21],
                                       k_h=3,
                                       k_w=3,
                                       d_h=1,
                                       d_w=1,
                                       padding=padding,
                                       initializer='He',
                                       name='s9_conv3')
            s9_conv3 = tf_utils.norm(s9_conv3,
                                     name='s9_norm2',
                                     _type='batch',
                                     _ops=self._ops,
                                     is_train=False)
            s9_conv3 = tf_utils.relu(s9_conv3, name='relu_s9_conv3')

            s10_deconv1 = tf_utils.deconv2d(x=s9_conv3,
                                            output_dim=self.conv_dims[-1],
                                            k_h=2,
                                            k_w=2,
                                            initializer='He',
                                            name='s10_deconv1')
            s10_deconv1 = tf_utils.norm(s10_deconv1,
                                        name='s10_norm0',
                                        _type='batch',
                                        _ops=self._ops,
                                        is_train=False)
            s10_deconv1 = tf_utils.relu(s10_deconv1, name='relu_s10_deconv1')
            # Concat
            s10_concat = tf_utils.concat(values=[s10_deconv1, s0_conv2],
                                         axis=3,
                                         name='s10_axis3_concat')

            s10_conv2 = tf_utils.conv2d(s10_concat,
                                        output_dim=self.conv_dims[-1],
                                        k_h=3,
                                        k_w=3,
                                        d_h=1,
                                        d_w=1,
                                        padding=padding,
                                        initializer='He',
                                        name='s10_conv2')
            s10_conv2 = tf_utils.norm(s10_conv2,
                                      name='s10_norm1',
                                      _type='batch',
                                      _ops=self._ops,
                                      is_train=False)
            s10_conv2 = tf_utils.relu(s10_conv2, name='relu_s10_conv2')

            s10_conv3 = tf_utils.conv2d(x=s10_conv2,
                                        output_dim=self.conv_dims[-1],
                                        k_h=3,
                                        k_w=3,
                                        d_h=1,
                                        d_w=1,
                                        padding=padding,
                                        initializer='He',
                                        name='s10_conv3')
            s10_conv3 = tf_utils.norm(s10_conv3,
                                      name='s10_norm2',
                                      _type='batch',
                                      _ops=self._ops,
                                      is_train=False)
            s10_conv3 = tf_utils.relu(s10_conv3, name='relu_s10_conv3')

            output = tf_utils.conv2d(s10_conv3,
                                     output_dim=self.num_classes,
                                     k_h=1,
                                     k_w=1,
                                     d_h=1,
                                     d_w=1,
                                     padding=padding,
                                     initializer='He',
                                     name='output')

            return output
    def __call__(self, x, keep_rate=0.5):
        with tf.compat.v1.variable_scope(self.name, reuse=self.reuse):
            tf_utils.print_activations(x, logger=self.logger)

            # E0: (320, 200) -> (160, 100)
            e0_conv2d = tf_utils.conv2d(x,
                                        output_dim=self.gen_c[0],
                                        initializer='He',
                                        logger=self.logger,
                                        name='e0_conv2d')
            e0_lrelu = tf_utils.lrelu(e0_conv2d,
                                      logger=self.logger,
                                      name='e0_lrelu')

            # E1: (160, 100) -> (80, 50)
            e1_conv2d = tf_utils.conv2d(e0_lrelu,
                                        output_dim=self.gen_c[1],
                                        initializer='He',
                                        logger=self.logger,
                                        name='e1_conv2d')
            e1_batchnorm = tf_utils.norm(e1_conv2d,
                                         _type=self.norm,
                                         _ops=self._ops,
                                         logger=self.logger,
                                         name='e1_norm')
            e1_lrelu = tf_utils.lrelu(e1_batchnorm,
                                      logger=self.logger,
                                      name='e1_lrelu')

            # E2: (80, 50) -> (40, 25)
            e2_conv2d = tf_utils.conv2d(e1_lrelu,
                                        output_dim=self.gen_c[2],
                                        initializer='He',
                                        logger=self.logger,
                                        name='e2_conv2d')
            e2_batchnorm = tf_utils.norm(e2_conv2d,
                                         _type=self.norm,
                                         _ops=self._ops,
                                         logger=self.logger,
                                         name='e2_norm')
            e2_lrelu = tf_utils.lrelu(e2_batchnorm,
                                      logger=self.logger,
                                      name='e2_lrelu')

            # E3: (40, 25) -> (20, 13)
            e3_conv2d = tf_utils.conv2d(e2_lrelu,
                                        output_dim=self.gen_c[3],
                                        initializer='He',
                                        logger=self.logger,
                                        name='e3_conv2d')
            e3_batchnorm = tf_utils.norm(e3_conv2d,
                                         _type=self.norm,
                                         _ops=self._ops,
                                         logger=self.logger,
                                         name='e3_norm')
            e3_lrelu = tf_utils.lrelu(e3_batchnorm,
                                      logger=self.logger,
                                      name='e3_lrelu')

            # E4: (20, 13) -> (10, 7)
            e4_conv2d = tf_utils.conv2d(e3_lrelu,
                                        output_dim=self.gen_c[4],
                                        initializer='He',
                                        logger=self.logger,
                                        name='e4_conv2d')
            e4_batchnorm = tf_utils.norm(e4_conv2d,
                                         _type=self.norm,
                                         _ops=self._ops,
                                         logger=self.logger,
                                         name='e4_norm')
            e4_lrelu = tf_utils.lrelu(e4_batchnorm,
                                      logger=self.logger,
                                      name='e4_lrelu')

            # E5: (10, 7) -> (5, 4)
            e5_conv2d = tf_utils.conv2d(e4_lrelu,
                                        output_dim=self.gen_c[5],
                                        initializer='He',
                                        logger=self.logger,
                                        name='e5_conv2d')
            e5_batchnorm = tf_utils.norm(e5_conv2d,
                                         _type=self.norm,
                                         _ops=self._ops,
                                         logger=self.logger,
                                         name='e5_norm')
            e5_lrelu = tf_utils.lrelu(e5_batchnorm,
                                      logger=self.logger,
                                      name='e5_lrelu')

            # E6: (5, 4) -> (3, 2)
            e6_conv2d = tf_utils.conv2d(e5_lrelu,
                                        output_dim=self.gen_c[6],
                                        initializer='He',
                                        logger=self.logger,
                                        name='e6_conv2d')
            e6_batchnorm = tf_utils.norm(e6_conv2d,
                                         _type=self.norm,
                                         _ops=self._ops,
                                         logger=self.logger,
                                         name='e6_norm')
            e6_lrelu = tf_utils.lrelu(e6_batchnorm,
                                      logger=self.logger,
                                      name='e6_lrelu')

            # E7: (3, 2) -> (2, 1)
            e7_conv2d = tf_utils.conv2d(e6_lrelu,
                                        output_dim=self.gen_c[7],
                                        initializer='He',
                                        logger=self.logger,
                                        name='e7_conv2d')
            e7_batchnorm = tf_utils.norm(e7_conv2d,
                                         _type=self.norm,
                                         _ops=self._ops,
                                         logger=self.logger,
                                         name='e7_norm')
            e7_relu = tf_utils.lrelu(e7_batchnorm,
                                     logger=self.logger,
                                     name='e7_relu')

            # D0: (2, 1) -> (3, 2)
            # Stage1: (2, 1) -> (4, 2)
            d0_deconv = tf_utils.deconv2d(e7_relu,
                                          output_dim=self.gen_c[8],
                                          initializer='He',
                                          logger=self.logger,
                                          name='d0_deconv2d')
            # Stage2: (4, 2) -> (3, 2)
            shapeA = e6_conv2d.get_shape().as_list()[1]
            shapeB = d0_deconv.get_shape().as_list()[1] - e6_conv2d.get_shape(
            ).as_list()[1]
            d0_split, _ = tf.split(d0_deconv, [shapeA, shapeB],
                                   axis=1,
                                   name='d0_split')
            tf_utils.print_activations(d0_split, logger=self.logger)
            # Stage3: Batch norm, concatenation, and relu
            d0_batchnorm = tf_utils.norm(d0_split,
                                         _type=self.norm,
                                         _ops=self._ops,
                                         logger=self.logger,
                                         name='d0_norm')
            d0_drop = tf_utils.dropout(d0_batchnorm,
                                       keep_prob=keep_rate,
                                       logger=self.logger,
                                       name='d0_dropout')
            d0_concat = tf.concat([d0_drop, e6_batchnorm],
                                  axis=3,
                                  name='d0_concat')
            d0_relu = tf_utils.relu(d0_concat,
                                    logger=self.logger,
                                    name='d0_relu')

            # D1: (3, 2) -> (5, 4)
            # Stage1: (3, 2) -> (6, 4)
            d1_deconv = tf_utils.deconv2d(d0_relu,
                                          output_dim=self.gen_c[9],
                                          initializer='He',
                                          logger=self.logger,
                                          name='d1_deconv2d')
            # Stage2: (6, 4) -> (5, 4)
            shapeA = e5_batchnorm.get_shape().as_list()[1]
            shapeB = d1_deconv.get_shape().as_list(
            )[1] - e5_batchnorm.get_shape().as_list()[1]
            d1_split, _ = tf.split(d1_deconv, [shapeA, shapeB],
                                   axis=1,
                                   name='d1_split')
            tf_utils.print_activations(d1_split, logger=self.logger)
            # Stage3: Batch norm, concatenation, and relu
            d1_batchnorm = tf_utils.norm(d1_split,
                                         _type=self.norm,
                                         _ops=self._ops,
                                         logger=self.logger,
                                         name='d1_norm')
            d1_drop = tf_utils.dropout(d1_batchnorm,
                                       keep_prob=keep_rate,
                                       logger=self.logger,
                                       name='d1_dropout')
            d1_concat = tf.concat([d1_drop, e5_batchnorm],
                                  axis=3,
                                  name='d1_concat')
            d1_relu = tf_utils.relu(d1_concat,
                                    logger=self.logger,
                                    name='d1_relu')

            # D2: (5, 4) -> (10, 7)
            # Stage1: (5, 4) -> (10, 8)
            d2_deconv = tf_utils.deconv2d(d1_relu,
                                          output_dim=self.gen_c[10],
                                          initializer='He',
                                          logger=self.logger,
                                          name='d2_deconv2d')
            # Stage2: (10, 8) -> (10, 7)
            shapeA = e4_batchnorm.get_shape().as_list()[2]
            shapeB = d2_deconv.get_shape().as_list(
            )[2] - e4_batchnorm.get_shape().as_list()[2]
            d2_split, _ = tf.split(d2_deconv, [shapeA, shapeB],
                                   axis=2,
                                   name='d2_split')
            tf_utils.print_activations(d2_split, logger=self.logger)
            # Stage3: Batch norm, concatenation, and relu
            d2_batchnorm = tf_utils.norm(d2_split,
                                         _type=self.norm,
                                         _ops=self._ops,
                                         logger=self.logger,
                                         name='d2_norm')
            d2_drop = tf_utils.dropout(d2_batchnorm,
                                       keep_prob=keep_rate,
                                       logger=self.logger,
                                       name='d2_dropout')
            d2_concat = tf.concat([d2_drop, e4_batchnorm],
                                  axis=3,
                                  name='d2_concat')
            d2_relu = tf_utils.relu(d2_concat,
                                    logger=self.logger,
                                    name='d2_relu')

            # D3: (10, 7) -> (20, 13)
            # Stage1: (10, 7) -> (20, 14)
            d3_deconv = tf_utils.deconv2d(d2_relu,
                                          output_dim=self.gen_c[11],
                                          initializer='He',
                                          logger=self.logger,
                                          name='d3_deconv2d')
            # Stage2: (20, 14) -> (20, 13)
            shapeA = e3_batchnorm.get_shape().as_list()[2]
            shapeB = d3_deconv.get_shape().as_list(
            )[2] - e3_batchnorm.get_shape().as_list()[2]
            d3_split, _ = tf.split(d3_deconv, [shapeA, shapeB],
                                   axis=2,
                                   name='d3_split_2')
            tf_utils.print_activations(d3_split, logger=self.logger)
            # Stage3: Batch norm, concatenation, and relu
            d3_batchnorm = tf_utils.norm(d3_split,
                                         _type=self.norm,
                                         _ops=self._ops,
                                         logger=self.logger,
                                         name='d3_norm')
            d3_concat = tf.concat([d3_batchnorm, e3_batchnorm],
                                  axis=3,
                                  name='d3_concat')
            d3_relu = tf_utils.relu(d3_concat,
                                    logger=self.logger,
                                    name='d3_relu')

            # D4: (20, 13) -> (40, 25)
            # Stage1: (20, 13) -> (40, 26)
            d4_deconv = tf_utils.deconv2d(d3_relu,
                                          output_dim=self.gen_c[12],
                                          initializer='He',
                                          logger=self.logger,
                                          name='d4_deconv2d')
            # Stage2: (40, 26) -> (40, 25)
            shapeA = e2_batchnorm.get_shape().as_list()[2]
            shapeB = d4_deconv.get_shape().as_list(
            )[2] - e2_batchnorm.get_shape().as_list()[2]
            d4_split, _ = tf.split(d4_deconv, [shapeA, shapeB],
                                   axis=2,
                                   name='d4_split')
            tf_utils.print_activations(d4_split, logger=self.logger)
            # Stage3: Batch norm, concatenation, and relu
            d4_batchnorm = tf_utils.norm(d4_split,
                                         _type=self.norm,
                                         _ops=self._ops,
                                         logger=self.logger,
                                         name='d4_norm')
            d4_concat = tf.concat([d4_batchnorm, e2_batchnorm],
                                  axis=3,
                                  name='d4_concat')
            d4_relu = tf_utils.relu(d4_concat,
                                    logger=self.logger,
                                    name='d4_relu')

            # D5: (40, 25, 256) -> (80, 50, 128)
            d5_deconv = tf_utils.deconv2d(d4_relu,
                                          output_dim=self.gen_c[13],
                                          initializer='He',
                                          logger=self.logger,
                                          name='d5_deconv2d')
            d5_batchnorm = tf_utils.norm(d5_deconv,
                                         _type=self.norm,
                                         _ops=self._ops,
                                         logger=self.logger,
                                         name='d5_norm')
            d5_concat = tf.concat([d5_batchnorm, e1_batchnorm],
                                  axis=3,
                                  name='d5_concat')
            d5_relu = tf_utils.relu(d5_concat,
                                    logger=self.logger,
                                    name='d5_relu')

            # D6: (80, 50, 128) -> (160, 100, 64)
            d6_deconv = tf_utils.deconv2d(d5_relu,
                                          output_dim=self.gen_c[14],
                                          initializer='He',
                                          logger=self.logger,
                                          name='d6_deconv2d')
            d6_batchnorm = tf_utils.norm(d6_deconv,
                                         _type=self.norm,
                                         _ops=self._ops,
                                         logger=self.logger,
                                         name='d6_norm')
            d6_concat = tf.concat([d6_batchnorm, e0_conv2d],
                                  axis=3,
                                  name='d6_concat')
            d6_relu = tf_utils.relu(d6_concat,
                                    logger=self.logger,
                                    name='d6_relu')

            # D7: (160, 100, 64) -> (320, 200, 1)
            d7_deconv = tf_utils.deconv2d(d6_relu,
                                          output_dim=self.gen_c[15],
                                          initializer='He',
                                          logger=self.logger,
                                          name='d7_deconv2d')
            output = tf_utils.tanh(d7_deconv,
                                   logger=self.logger,
                                   name='output_tanh')

            # Set reuse=True for next call
            self.reuse = True
            self.variables = tf.compat.v1.get_collection(
                tf.compat.v1.GraphKeys.TRAINABLE_VARIABLES, scope=self.name)

        return output
    def forward_network(self, inputImg, padding='SAME', reuse=False):
        with tf.compat.v1.variable_scope(self.name, reuse=reuse):
            # This part is for compatible between input size [640, 400] and [320, 200]
            if self.resize_factor == 1.0:
                # Stage 0
                tf_utils.print_activations(inputImg, logger=self.logger)
                s0_conv1 = tf_utils.conv2d(x=inputImg, output_dim=self.conv_dims[0], k_h=3, k_w=3, d_h=1, d_w=1,
                                           padding=padding, initializer='He', name='s0_conv1', logger=self.logger)
                s0_conv1 = tf_utils.relu(s0_conv1, name='relu_s0_conv1', logger=self.logger)

                s0_conv2 = tf_utils.conv2d(x=s0_conv1, output_dim=self.conv_dims[0], k_h=3, k_w=3, d_h=1, d_w=1,
                                           padding=padding, initializer='He', name='s0_conv2', logger=self.logger)
                if self.use_batch_norm:
                    s0_conv2 = tf_utils.norm(s0_conv2, name='s0_norm1', _type='batch', _ops=self._ops,
                                             is_train=self.trainMode, logger=self.logger)
                s0_conv2 = tf_utils.relu(s0_conv2, name='relu_s0_conv2', logger=self.logger)

                # Stage 1
                s1_maxpool = tf_utils.max_pool(x=s0_conv2, name='s1_maxpool2d', logger=self.logger)

                s1_conv1 = tf_utils.conv2d(x=s1_maxpool, output_dim=self.conv_dims[0], k_h=3, k_w=3, d_h=1, d_w=1,
                                           padding=padding, initializer='He', name='s1_conv1', logger=self.logger)
                if self.use_batch_norm:
                    s1_conv1 = tf_utils.norm(s1_conv1, name='s1_norm0', _type='batch', _ops=self._ops,
                                             is_train=self.trainMode, logger=self.logger)
                s1_conv1 = tf_utils.relu(s1_conv1, name='relu_s1_conv1', logger=self.logger)

                s1_conv2 = tf_utils.conv2d(x=s1_conv1, output_dim=self.conv_dims[1], k_h=3, k_w=3, d_h=1, d_w=1,
                                           padding=padding, initializer='He', name='s1_conv2', logger=self.logger)
                if self.use_batch_norm:
                    s1_conv2 = tf_utils.norm(s1_conv2, name='s1_norm1', _type='batch', _ops=self._ops,
                                             is_train=self.trainMode, logger=self.logger)
                s1_conv2 = tf_utils.relu(s1_conv2, name='relu_s1_conv2', logger=self.logger)
            else:
                # Stage 1
                tf_utils.print_activations(inputImg, logger=self.logger)
                s1_conv1 = tf_utils.conv2d(x=inputImg, output_dim=self.conv_dims[0], k_h=3, k_w=3, d_h=1, d_w=1,
                                           padding=padding, initializer='He', name='s1_conv1', logger=self.logger)
                s1_conv1 = tf_utils.relu(s1_conv1, name='relu_s1_conv1', logger=self.logger)

                s1_conv2 = tf_utils.conv2d(x=s1_conv1, output_dim=self.conv_dims[1], k_h=3, k_w=3, d_h=1, d_w=1,
                                           padding=padding, initializer='He', name='s1_conv2', logger=self.logger)
                if self.use_batch_norm:
                    s1_conv2 = tf_utils.norm(s1_conv2, name='s1_norm1', _type='batch', _ops=self._ops,
                                             is_train=self.trainMode, logger=self.logger)
                s1_conv2 = tf_utils.relu(s1_conv2, name='relu_s1_conv2', logger=self.logger)

            # Stage 2
            s2_maxpool = tf_utils.max_pool(x=s1_conv2, name='s2_maxpool2d', logger=self.logger)
            s2_conv1 = tf_utils.conv2d(x=s2_maxpool, output_dim=self.conv_dims[2], k_h=3, k_w=3, d_h=1, d_w=1,
                                       padding=padding, initializer='He', name='s2_conv1', logger=self.logger)
            if self.use_batch_norm:
                s2_conv1 = tf_utils.norm(s2_conv1, name='s2_norm0', _type='batch', _ops=self._ops,
                                         is_train=self.trainMode, logger=self.logger)
            s2_conv1 = tf_utils.relu(s2_conv1, name='relu_s2_conv1', logger=self.logger)

            s2_conv2 = tf_utils.conv2d(x=s2_conv1, output_dim=self.conv_dims[3], k_h=3, k_w=3, d_h=1, d_w=1,
                                       padding=padding, initializer='He', name='s2_conv2', logger=self.logger)
            if self.use_batch_norm:
                s2_conv2 = tf_utils.norm(s2_conv2, name='s2_norm1', _type='batch', _ops=self._ops,
                                         is_train=self.trainMode, logger=self.logger)
            s2_conv2 = tf_utils.relu(s2_conv2, name='relu_s2_conv2', logger=self.logger)

            # Stage 3
            s3_maxpool = tf_utils.max_pool(x=s2_conv2, name='s3_maxpool2d', logger=self.logger)
            s3_conv1 = tf_utils.conv2d(x=s3_maxpool, output_dim=self.conv_dims[4], k_h=3, k_w=3, d_h=1, d_w=1,
                                       padding=padding, initializer='He', name='s3_conv1', logger=self.logger)
            if self.use_batch_norm:
                s3_conv1 = tf_utils.norm(s3_conv1, name='s3_norm0', _type='batch', _ops=self._ops,
                                         is_train=self.trainMode, logger=self.logger)
            s3_conv1 = tf_utils.relu(s3_conv1, name='relu_s3_conv1', logger=self.logger)

            s3_conv2 = tf_utils.conv2d(x=s3_conv1, output_dim=self.conv_dims[5], k_h=3, k_w=3, d_h=1, d_w=1,
                                       padding=padding, initializer='He', name='s3_conv2', logger=self.logger)
            if self.use_batch_norm:
                s3_conv2 = tf_utils.norm(s3_conv2, name='s3_norm1', _type='batch', _ops=self._ops,
                                         is_train=self.trainMode, logger=self.logger)
            s3_conv2 = tf_utils.relu(s3_conv2, name='relu_s3_conv2', logger=self.logger)

            # Stage 4
            s4_maxpool = tf_utils.max_pool(x=s3_conv2, name='s4_maxpool2d', logger=self.logger)
            s4_conv1 = tf_utils.conv2d(x=s4_maxpool, output_dim=self.conv_dims[6], k_h=3, k_w=3, d_h=1, d_w=1,
                                       padding=padding, initializer='He', name='s4_conv1', logger=self.logger)
            if self.use_batch_norm:
                s4_conv1 = tf_utils.norm(s4_conv1, name='s4_norm0', _type='batch', _ops=self._ops,
                                         is_train=self.trainMode, logger=self.logger)
            s4_conv1 = tf_utils.relu(s4_conv1, name='relu_s4_conv1', logger=self.logger)

            s4_conv2 = tf_utils.conv2d(x=s4_conv1, output_dim=self.conv_dims[7], k_h=3, k_w=3, d_h=1, d_w=1,
                                       padding=padding, initializer='He', name='s4_conv2', logger=self.logger)
            if self.use_batch_norm:
                s4_conv2 = tf_utils.norm(s4_conv2, name='s4_norm1', _type='batch', _ops=self._ops,
                                         is_train=self.trainMode, logger=self.logger)
            s4_conv2 = tf_utils.relu(s4_conv2, name='relu_s4_conv2', logger=self.logger)
            s4_conv2_drop = tf_utils.dropout(x=s4_conv2, keep_prob=self.ratePh, name='s4_dropout',
                                             logger=self.logger)

            # Stage 5
            s5_maxpool = tf_utils.max_pool(x=s4_conv2_drop, name='s5_maxpool2d', logger=self.logger)
            s5_conv1 = tf_utils.conv2d(x=s5_maxpool, output_dim=self.conv_dims[8], k_h=3, k_w=3, d_h=1, d_w=1,
                                       padding=padding, initializer='He', name='s5_conv1', logger=self.logger)
            if self.use_batch_norm:
                s5_conv1 = tf_utils.norm(s5_conv1, name='s5_norm0', _type='batch', _ops=self._ops,
                                         is_train=self.trainMode, logger=self.logger)
            s5_conv1 = tf_utils.relu(s5_conv1, name='relu_s5_conv1', logger=self.logger)

            s5_conv2 = tf_utils.conv2d(x=s5_conv1, output_dim=self.conv_dims[9], k_h=3, k_w=3, d_h=1, d_w=1,
                                       padding=padding, initializer='He', name='s5_conv2', logger=self.logger)
            if self.use_batch_norm:
                s5_conv2 = tf_utils.norm(s5_conv2, name='s5_norm1', _type='batch', _ops=self._ops,
                                         is_train=self.trainMode, logger=self.logger)
            s5_conv2 = tf_utils.relu(s5_conv2, name='relu_s5_conv2', logger=self.logger)
            s5_conv2_drop = tf_utils.dropout(x=s5_conv2, keep_prob=self.ratePh, name='s5_dropout',
                                             logger=self.logger)

            # Stage 6
            s6_deconv1 = tf_utils.deconv2d(x=s5_conv2_drop, output_dim=self.conv_dims[10], k_h=2, k_w=2,
                                           initializer='He', name='s6_deconv1', logger=self.logger)
            if self.use_batch_norm:
                s6_deconv1 = tf_utils.norm(s6_deconv1, name='s6_norm0', _type='batch', _ops=self._ops,
                                         is_train=self.trainMode, logger=self.logger)
            s6_deconv1 = tf_utils.relu(s6_deconv1, name='relu_s6_deconv1', logger=self.logger)
            # Cropping
            w1 = s4_conv2_drop.get_shape().as_list()[2]
            w2 = s6_deconv1.get_shape().as_list()[2] - s4_conv2_drop.get_shape().as_list()[2]
            s6_deconv1_split, _ = tf.split(s6_deconv1, num_or_size_splits=[w1, w2], axis=2, name='axis2_split')
            tf_utils.print_activations(s6_deconv1_split, logger=self.logger)
            # Concat
            s6_concat = tf_utils.concat(values=[s6_deconv1_split, s4_conv2_drop], axis=3, name='s6_axis3_concat',
                                        logger=self.logger)

            s6_conv2 = tf_utils.conv2d(x=s6_concat, output_dim=self.conv_dims[11], k_h=3, k_w=3, d_h=1, d_w=1,
                                       padding=padding, initializer='He', name='s6_conv2', logger=self.logger)
            if self.use_batch_norm:
                s6_conv2 = tf_utils.norm(s6_conv2, name='s6_norm1', _type='batch', _ops=self._ops,
                                         is_train=self.trainMode, logger=self.logger)
            s6_conv2 = tf_utils.relu(s6_conv2, name='relu_s6_conv2', logger=self.logger)

            s6_conv3 = tf_utils.conv2d(x=s6_conv2, output_dim=self.conv_dims[12], k_h=3, k_w=3, d_h=1, d_w=1,
                                       padding=padding, initializer='He', name='s6_conv3', logger=self.logger)
            if self.use_batch_norm:
                s6_conv3 = tf_utils.norm(s6_conv3, name='s6_norm2', _type='batch', _ops=self._ops,
                                         is_train=self.trainMode, logger=self.logger)
            s6_conv3 = tf_utils.relu(s6_conv3, name='relu_s6_conv3', logger=self.logger)

            # Stage 7
            s7_deconv1 = tf_utils.deconv2d(x=s6_conv3, output_dim=self.conv_dims[13], k_h=2, k_w=2, initializer='He',
                                           name='s7_deconv1', logger=self.logger)
            if self.use_batch_norm:
                s7_deconv1 = tf_utils.norm(s7_deconv1, name='s7_norm0', _type='batch', _ops=self._ops,
                                         is_train=self.trainMode, logger=self.logger)
            s7_deconv1 = tf_utils.relu(s7_deconv1, name='relu_s7_deconv1', logger=self.logger)
            # Concat
            s7_concat = tf_utils.concat(values=[s7_deconv1, s3_conv2], axis=3, name='s7_axis3_concat',
                                        logger=self.logger)

            s7_conv2 = tf_utils.conv2d(x=s7_concat, output_dim=self.conv_dims[14], k_h=3, k_w=3, d_h=1, d_w=1,
                                       padding=padding, initializer='He', name='s7_conv2', logger=self.logger)
            if self.use_batch_norm:
                s7_conv2 = tf_utils.norm(s7_conv2, name='s7_norm1', _type='batch', _ops=self._ops,
                                         is_train=self.trainMode, logger=self.logger)
            s7_conv2 = tf_utils.relu(s7_conv2, name='relu_s7_conv2', logger=self.logger)

            s7_conv3 = tf_utils.conv2d(x=s7_conv2, output_dim=self.conv_dims[15], k_h=3, k_w=3, d_h=1, d_w=1,
                                       padding=padding, initializer='He', name='s7_conv3', logger=self.logger)
            if self.use_batch_norm:
                s7_conv3 = tf_utils.norm(s7_conv3, name='s7_norm2', _type='batch', _ops=self._ops,
                                         is_train=self.trainMode, logger=self.logger)
            s7_conv3 = tf_utils.relu(s7_conv3, name='relu_s7_conv3', logger=self.logger)

            # Stage 8
            s8_deconv1 = tf_utils.deconv2d(x=s7_conv3, output_dim=self.conv_dims[16], k_h=2, k_w=2, initializer='He',
                                           name='s8_deconv1', logger=self.logger)
            if self.use_batch_norm:
                s8_deconv1 = tf_utils.norm(s8_deconv1, name='s8_norm0', _type='batch', _ops=self._ops,
                                         is_train=self.trainMode, logger=self.logger)
            s8_deconv1 = tf_utils.relu(s8_deconv1, name='relu_s8_deconv1', logger=self.logger)
            # Concat
            s8_concat = tf_utils.concat(values=[s8_deconv1,s2_conv2], axis=3, name='s8_axis3_concat',
                                        logger=self.logger)

            s8_conv2 = tf_utils.conv2d(x=s8_concat, output_dim=self.conv_dims[17], k_h=3, k_w=3, d_h=1, d_w=1,
                                       padding=padding, initializer='He', name='s8_conv2', logger=self.logger)
            if self.use_batch_norm:
                s8_conv2 = tf_utils.norm(s8_conv2, name='s8_norm1', _type='batch', _ops=self._ops,
                                         is_train=self.trainMode, logger=self.logger)
            s8_conv2 = tf_utils.relu(s8_conv2, name='relu_s8_conv2', logger=self.logger)

            s8_conv3 = tf_utils.conv2d(x=s8_conv2, output_dim=self.conv_dims[18], k_h=3, k_w=3, d_h=1, d_w=1,
                                       padding=padding, initializer='He', name='s8_conv3', logger=self.logger)
            if self.use_batch_norm:
                s8_conv3 = tf_utils.norm(s8_conv3, name='s8_norm2', _type='batch', _ops=self._ops,
                                         is_train=self.trainMode, logger=self.logger)
            s8_conv3 = tf_utils.relu(s8_conv3, name='relu_conv3', logger=self.logger)

            # Stage 9
            s9_deconv1 = tf_utils.deconv2d(x=s8_conv3, output_dim=self.conv_dims[19], k_h=2, k_w=2,
                                           initializer='He', name='s9_deconv1', logger=self.logger)
            if self.use_batch_norm:
                s9_deconv1 = tf_utils.norm(s9_deconv1, name='s9_norm0', _type='batch', _ops=self._ops,
                                         is_train=self.trainMode, logger=self.logger)
            s9_deconv1 = tf_utils.relu(s9_deconv1, name='relu_s9_deconv1', logger=self.logger)
            # Concat
            s9_concat = tf_utils.concat(values=[s9_deconv1, s1_conv2], axis=3, name='s9_axis3_concat',
                                        logger=self.logger)

            s9_conv2 = tf_utils.conv2d(x=s9_concat, output_dim=self.conv_dims[20], k_h=3, k_w=3, d_h=1, d_w=1,
                                       padding=padding, initializer='He', name='s9_conv2', logger=self.logger)
            if self.use_batch_norm:
                s9_conv2 = tf_utils.norm(s9_conv2, name='s9_norm1', _type='batch', _ops=self._ops,
                                         is_train=self.trainMode, logger=self.logger)
            s9_conv2 = tf_utils.relu(s9_conv2, name='relu_s9_conv2', logger=self.logger)

            s9_conv3 = tf_utils.conv2d(x=s9_conv2, output_dim=self.conv_dims[21], k_h=3, k_w=3, d_h=1, d_w=1,
                                       padding=padding, initializer='He', name='s9_conv3', logger=self.logger)
            if self.use_batch_norm:
                s9_conv3 = tf_utils.norm(s9_conv3, name='s9_norm2', _type='batch', _ops=self._ops,
                                         is_train=self.trainMode, logger=self.logger)
            s9_conv3 = tf_utils.relu(s9_conv3, name='relu_s9_conv3', logger=self.logger)

            if self.resize_factor == 1.0:
                s10_deconv1 = tf_utils.deconv2d(x=s9_conv3, output_dim=self.conv_dims[-1], k_h=2, k_w=2,
                                                initializer='He', name='s10_deconv1', logger=self.logger)
                if self.use_batch_norm:
                    s10_deconv1 = tf_utils.norm(s10_deconv1, name='s10_norm0', _type='batch', _ops=self._ops,
                                             is_train=self.trainMode, logger=self.logger)
                s10_deconv1 = tf_utils.relu(s10_deconv1, name='relu_s10_deconv1', logger=self.logger)
                # Concat
                s10_concat = tf_utils.concat(values=[s10_deconv1, s0_conv2], axis=3, name='s10_axis3_concat',
                                             logger=self.logger)

                s10_conv2 = tf_utils.conv2d(s10_concat, output_dim=self.conv_dims[-1], k_h=3, k_w=3, d_h=1, d_w=1,
                                            padding=padding, initializer='He', name='s10_conv2', logger=self.logger)
                if self.use_batch_norm:
                    s10_conv2 = tf_utils.norm(s10_conv2, name='s10_norm1', _type='batch', _ops=self._ops,
                                             is_train=self.trainMode, logger=self.logger)
                s10_conv2 = tf_utils.relu(s10_conv2, name='relu_s10_conv2', logger=self.logger)

                s10_conv3 = tf_utils.conv2d(x=s10_conv2, output_dim=self.conv_dims[-1], k_h=3, k_w=3, d_h=1, d_w=1,
                                            padding=padding, initializer='He', name='s10_conv3', logger=self.logger)
                if self.use_batch_norm:
                    s10_conv3 = tf_utils.norm(s10_conv3, name='s10_norm2', _type='batch', _ops=self._ops,
                                             is_train=self.trainMode, logger=self.logger)
                s10_conv3 = tf_utils.relu(s10_conv3, name='relu_s10_conv3', logger=self.logger)

                output = tf_utils.conv2d(s10_conv3, output_dim=self.numClasses, k_h=1, k_w=1, d_h=1, d_w=1,
                                         padding=padding, initializer='He', name='output', logger=self.logger)
            else:
                output = tf_utils.conv2d(s9_conv3, output_dim=self.numClasses, k_h=1, k_w=1, d_h=1, d_w=1,
                                         padding=padding, initializer='He', name='output', logger=self.logger)

            return output
Пример #5
0
    def u_net(self):
        # Stage 1
        tf_utils.print_activations(self.inp_img, logger=self.logger)
        s1_conv1 = tf_utils.conv2d(x=self.inp_img,
                                   output_dim=self.conv_dims[0],
                                   k_h=3,
                                   k_w=3,
                                   d_h=1,
                                   d_w=1,
                                   padding='VALID',
                                   initializer='He',
                                   name='s1_conv1',
                                   logger=self.logger)
        s1_conv1 = tf_utils.relu(s1_conv1,
                                 name='relu_s1_conv1',
                                 logger=self.logger)
        s1_conv2 = tf_utils.conv2d(x=s1_conv1,
                                   output_dim=self.conv_dims[1],
                                   k_h=3,
                                   k_w=3,
                                   d_h=1,
                                   d_w=1,
                                   padding='VALID',
                                   initializer='He',
                                   name='s1_conv2',
                                   logger=self.logger)
        s1_conv2 = tf_utils.relu(s1_conv2,
                                 name='relu_s1_conv2',
                                 logger=self.logger)

        # Stage 2
        s2_maxpool = tf_utils.max_pool(x=s1_conv2,
                                       name='s2_maxpool',
                                       logger=self.logger)
        s2_conv1 = tf_utils.conv2d(x=s2_maxpool,
                                   output_dim=self.conv_dims[2],
                                   k_h=3,
                                   k_w=3,
                                   d_h=1,
                                   d_w=1,
                                   padding='VALID',
                                   initializer='He',
                                   name='s2_conv1',
                                   logger=self.logger)
        s2_conv1 = tf_utils.relu(s2_conv1,
                                 name='relu_s2_conv1',
                                 logger=self.logger)
        s2_conv2 = tf_utils.conv2d(x=s2_conv1,
                                   output_dim=self.conv_dims[3],
                                   k_h=3,
                                   k_w=3,
                                   d_h=1,
                                   d_w=1,
                                   padding='VALID',
                                   initializer='He',
                                   name='s2_conv2',
                                   logger=self.logger)
        s2_conv2 = tf_utils.relu(s2_conv2,
                                 name='relu_s2_conv2',
                                 logger=self.logger)

        # Stage 3
        s3_maxpool = tf_utils.max_pool(x=s2_conv2,
                                       name='s3_maxpool',
                                       logger=self.logger)
        s3_conv1 = tf_utils.conv2d(x=s3_maxpool,
                                   output_dim=self.conv_dims[4],
                                   k_h=3,
                                   k_w=3,
                                   d_h=1,
                                   d_w=1,
                                   padding='VALID',
                                   initializer='He',
                                   name='s3_conv1',
                                   logger=self.logger)
        s3_conv1 = tf_utils.relu(s3_conv1,
                                 name='relu_s3_conv1',
                                 logger=self.logger)
        s3_conv2 = tf_utils.conv2d(x=s3_conv1,
                                   output_dim=self.conv_dims[5],
                                   k_h=3,
                                   k_w=3,
                                   d_h=1,
                                   d_w=1,
                                   padding='VALID',
                                   initializer='He',
                                   name='s3_conv2',
                                   logger=self.logger)
        s3_conv2 = tf_utils.relu(s3_conv2,
                                 name='relu_s3_conv2',
                                 logger=self.logger)

        # Stage 4
        s4_maxpool = tf_utils.max_pool(x=s3_conv2,
                                       name='s4_maxpool',
                                       logger=self.logger)
        s4_conv1 = tf_utils.conv2d(x=s4_maxpool,
                                   output_dim=self.conv_dims[6],
                                   k_h=3,
                                   k_w=3,
                                   d_h=1,
                                   d_w=1,
                                   padding='VALID',
                                   initializer='He',
                                   name='s4_conv1',
                                   logger=self.logger)
        s4_conv1 = tf_utils.relu(s4_conv1,
                                 name='relu_s4_conv1',
                                 logger=self.logger)
        s4_conv2 = tf_utils.conv2d(x=s4_conv1,
                                   output_dim=self.conv_dims[7],
                                   k_h=3,
                                   k_w=3,
                                   d_h=1,
                                   d_w=1,
                                   padding='VALID',
                                   initializer='He',
                                   name='s4_conv2',
                                   logger=self.logger)
        s4_conv2 = tf_utils.relu(s4_conv2,
                                 name='relu_s4_conv2',
                                 logger=self.logger)
        s4_conv2_drop = tf_utils.dropout(x=s4_conv2,
                                         keep_prob=self.keep_prob,
                                         name='s4_conv2_dropout',
                                         logger=self.logger)

        # Stage 5
        s5_maxpool = tf_utils.max_pool(x=s4_conv2_drop,
                                       name='s5_maxpool',
                                       logger=self.logger)
        s5_conv1 = tf_utils.conv2d(x=s5_maxpool,
                                   output_dim=self.conv_dims[8],
                                   k_h=3,
                                   k_w=3,
                                   d_h=1,
                                   d_w=1,
                                   padding='VALID',
                                   initializer='He',
                                   name='s5_conv1',
                                   logger=self.logger)
        s5_conv1 = tf_utils.relu(s5_conv1,
                                 name='relu_s5_conv1',
                                 logger=self.logger)
        s5_conv2 = tf_utils.conv2d(x=s5_conv1,
                                   output_dim=self.conv_dims[9],
                                   k_h=3,
                                   k_w=3,
                                   d_h=1,
                                   d_w=1,
                                   padding='VALID',
                                   initializer='He',
                                   name='s5_conv2',
                                   logger=self.logger)
        s5_conv2 = tf_utils.relu(s5_conv2,
                                 name='relu_s5_conv2',
                                 logger=self.logger)
        s5_conv2_drop = tf_utils.dropout(x=s5_conv2,
                                         keep_prob=self.keep_prob,
                                         name='s5_conv2_dropout',
                                         logger=self.logger)

        # Stage 6
        s6_deconv1 = tf_utils.deconv2d(x=s5_conv2_drop,
                                       output_dim=self.conv_dims[10],
                                       k_h=2,
                                       k_w=2,
                                       initializer='He',
                                       name='s6_deconv1',
                                       logger=self.logger)
        s6_deconv1 = tf_utils.relu(s6_deconv1,
                                   name='relu_s6_deconv1',
                                   logger=self.logger)

        # Cropping
        h1, w1 = s4_conv2_drop.get_shape().as_list()[1:3]
        h2, w2 = s6_deconv1.get_shape().as_list()[1:3]
        s4_conv2_crop = tf.image.crop_to_bounding_box(
            image=s4_conv2_drop,
            offset_height=int(0.5 * (h1 - h2)),
            offset_width=int(0.5 * (w1 - w2)),
            target_height=h2,
            target_width=w2)
        tf_utils.print_activations(s4_conv2_crop, logger=self.logger)

        s6_concat = tf_utils.concat(values=[s4_conv2_crop, s6_deconv1],
                                    axis=3,
                                    name='s6_concat',
                                    logger=self.logger)
        s6_conv2 = tf_utils.conv2d(x=s6_concat,
                                   output_dim=self.conv_dims[11],
                                   k_h=3,
                                   k_w=3,
                                   d_h=1,
                                   d_w=1,
                                   padding='VALID',
                                   initializer='He',
                                   name='s6_conv2',
                                   logger=self.logger)
        s6_conv2 = tf_utils.relu(s6_conv2,
                                 name='relu_s6_conv2',
                                 logger=self.logger)
        s6_conv3 = tf_utils.conv2d(x=s6_conv2,
                                   output_dim=self.conv_dims[12],
                                   k_h=3,
                                   k_w=3,
                                   d_h=1,
                                   d_w=1,
                                   padding='VALID',
                                   initializer='He',
                                   name='s6_conv3',
                                   logger=self.logger)
        s6_conv3 = tf_utils.relu(s6_conv3,
                                 name='relu_s6_conv3',
                                 logger=self.logger)

        # Stage 7
        s7_deconv1 = tf_utils.deconv2d(x=s6_conv3,
                                       output_dim=self.conv_dims[13],
                                       k_h=2,
                                       k_w=2,
                                       initializer='He',
                                       name='s7_deconv1',
                                       logger=self.logger)
        s7_deconv1 = tf_utils.relu(s7_deconv1,
                                   name='relu_s7_deconv1',
                                   logger=self.logger)
        # Cropping
        h1, w1 = s3_conv2.get_shape().as_list()[1:3]
        h2, w2 = s7_deconv1.get_shape().as_list()[1:3]
        s3_conv2_crop = tf.image.crop_to_bounding_box(
            image=s3_conv2,
            offset_height=int(0.5 * (h1 - h2)),
            offset_width=int(0.5 * (w1 - w2)),
            target_height=h2,
            target_width=w2)
        tf_utils.print_activations(s3_conv2_crop, logger=self.logger)

        s7_concat = tf_utils.concat(values=[s3_conv2_crop, s7_deconv1],
                                    axis=3,
                                    name='s7_concat',
                                    logger=self.logger)
        s7_conv2 = tf_utils.conv2d(x=s7_concat,
                                   output_dim=self.conv_dims[14],
                                   k_h=3,
                                   k_w=3,
                                   d_h=1,
                                   d_w=1,
                                   padding='VALID',
                                   initializer='He',
                                   name='s7_conv2',
                                   logger=self.logger)
        s7_conv2 = tf_utils.relu(s7_conv2,
                                 name='relu_s7_conv2',
                                 logger=self.logger)
        s7_conv3 = tf_utils.conv2d(x=s7_conv2,
                                   output_dim=self.conv_dims[15],
                                   k_h=3,
                                   k_w=3,
                                   d_h=1,
                                   d_w=1,
                                   padding='VALID',
                                   initializer='He',
                                   name='s7_conv3',
                                   logger=self.logger)
        s7_conv3 = tf_utils.relu(s7_conv3,
                                 name='relu_s7_conv3',
                                 logger=self.logger)

        # Stage 8
        s8_deconv1 = tf_utils.deconv2d(x=s7_conv3,
                                       output_dim=self.conv_dims[16],
                                       k_h=2,
                                       k_w=2,
                                       initializer='He',
                                       name='s8_deconv1',
                                       logger=self.logger)
        s8_deconv1 = tf_utils.relu(s8_deconv1,
                                   name='relu_s8_deconv1',
                                   logger=self.logger)
        # Cropping
        h1, w1 = s2_conv2.get_shape().as_list()[1:3]
        h2, w2 = s8_deconv1.get_shape().as_list()[1:3]
        s2_conv2_crop = tf.image.crop_to_bounding_box(
            image=s2_conv2,
            offset_height=int(0.5 * (h1 - h2)),
            offset_width=int(0.5 * (w1 - w2)),
            target_height=h2,
            target_width=w2)
        tf_utils.print_activations(s2_conv2_crop, logger=self.logger)

        s8_concat = tf_utils.concat(values=[s2_conv2_crop, s8_deconv1],
                                    axis=3,
                                    name='s8_concat',
                                    logger=self.logger)
        s8_conv2 = tf_utils.conv2d(x=s8_concat,
                                   output_dim=self.conv_dims[17],
                                   k_h=3,
                                   k_w=3,
                                   d_h=1,
                                   d_w=1,
                                   padding='VALID',
                                   initializer='He',
                                   name='s8_conv2',
                                   logger=self.logger)
        s8_conv2 = tf_utils.relu(s8_conv2,
                                 name='relu_s8_conv2',
                                 logger=self.logger)
        s8_conv3 = tf_utils.conv2d(x=s8_conv2,
                                   output_dim=self.conv_dims[18],
                                   k_h=3,
                                   k_w=3,
                                   d_h=1,
                                   d_w=1,
                                   padding='VALID',
                                   initializer='He',
                                   name='s8_conv3',
                                   logger=self.logger)
        s8_conv3 = tf_utils.relu(s8_conv3,
                                 name='relu_conv3',
                                 logger=self.logger)

        # Stage 9
        s9_deconv1 = tf_utils.deconv2d(x=s8_conv3,
                                       output_dim=self.conv_dims[19],
                                       k_h=2,
                                       k_w=2,
                                       initializer='He',
                                       name='s9_deconv1',
                                       logger=self.logger)
        s9_deconv1 = tf_utils.relu(s9_deconv1,
                                   name='relu_s9_deconv1',
                                   logger=self.logger)
        # Cropping
        h1, w1 = s1_conv2.get_shape().as_list()[1:3]
        h2, w2 = s9_deconv1.get_shape().as_list()[1:3]
        s1_conv2_crop = tf.image.crop_to_bounding_box(
            image=s1_conv2,
            offset_height=int(0.5 * (h1 - h2)),
            offset_width=int(0.5 * (w1 - w2)),
            target_height=h2,
            target_width=w2)
        tf_utils.print_activations(s1_conv2_crop, logger=self.logger)

        s9_concat = tf_utils.concat(values=[s1_conv2_crop, s9_deconv1],
                                    axis=3,
                                    name='s9_concat',
                                    logger=self.logger)
        s9_conv2 = tf_utils.conv2d(x=s9_concat,
                                   output_dim=self.conv_dims[20],
                                   k_h=3,
                                   k_w=3,
                                   d_h=1,
                                   d_w=1,
                                   padding='VALID',
                                   initializer='He',
                                   name='s9_conv2',
                                   logger=self.logger)
        s9_conv2 = tf_utils.relu(s9_conv2,
                                 name='relu_s9_conv2',
                                 logger=self.logger)
        s9_conv3 = tf_utils.conv2d(x=s9_conv2,
                                   output_dim=self.conv_dims[21],
                                   k_h=3,
                                   k_w=3,
                                   d_h=1,
                                   d_w=1,
                                   padding='VALID',
                                   initializer='He',
                                   name='s9_conv3',
                                   logger=self.logger)
        s9_conv3 = tf_utils.relu(s9_conv3,
                                 name='relu_s9_conv3',
                                 logger=self.logger)
        self.pred = tf_utils.conv2d(x=s9_conv3,
                                    output_dim=self.conv_dims[22],
                                    k_h=1,
                                    k_w=1,
                                    d_h=1,
                                    d_w=1,
                                    padding='SAME',
                                    initializer='He',
                                    name='output',
                                    logger=self.logger)
    def __init__(self, input_dim, output_dim=[1000, 1000, 10], optimizer=None, use_dropout=True, lr=0.001,
                 weight_decay=1e-4, random_seed=123, is_train=True, log_dir=None, name=None):
        self.name = name
        self.is_train = is_train
        self.log_dir = log_dir
        self.cur_lr = None
        self.logger, self.file_handler, self.stream_handler = utils.init_logger(log_dir=self.log_dir,
                                                                                name=self.name,
                                                                                is_train=self.is_train)

        with tf.variable_scope(self.name):
            # Placeholders for inputs
            self.X = tf.placeholder(dtype=tf.float32, shape=[None, input_dim], name='X')
            tf_utils.print_activations(self.X, logger=self.logger if self.is_train else None)
            self.y = tf.placeholder(dtype=tf.float32, shape=[None, output_dim[-1]], name='y')
            self.y_cls = tf.math.argmax(input=self.y, axis=1)
            self.keep_prob = tf.placeholder(tf.float32, name='keep_prob')

            # Placeholders for TensorBoard
            self.train_acc = tf.placeholder(tf.float32, name='train_acc')
            self.val_acc = tf.placeholder(tf.float32, name='val_acc')

            net = self.X
            for idx in range(len(output_dim) - 1):
                net = tf_utils.linear(x=net,
                                      output_size=output_dim[idx],
                                      name='fc'+str(idx),
                                      logger=self.logger if self.is_train else None)

                if use_dropout:
                    net = tf_utils.dropout(x=net,
                                           keep_prob=self.keep_prob,
                                           seed=random_seed,
                                           name='dropout'+str(idx),
                                           logger=self.logger if self.is_train else None)

                net = tf_utils.relu(x=net,
                                    name='relu'+str(idx),
                                    logger=self.logger if self.is_train else None)

            # Last predict layer
            self.y_pred = tf_utils.linear(net, output_size=output_dim[-1], name='last_fc')
            tf_utils.print_activations(self.y_pred, logger=self.logger if self.is_train else None)

            # Loss = data loss + regularization term
            self.data_loss = tf.math.reduce_mean(
                tf.nn.sigmoid_cross_entropy_with_logits(logits=self.y_pred, labels=self.y))
            self.reg_term = weight_decay * tf.reduce_sum(
                [tf.nn.l2_loss(weight) for weight in tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES)])
            self.loss = self.data_loss + self.reg_term

            # Optimizer
            self.train_op, self.cur_lr = optimizer_fn(optimizer, lr=lr, loss=self.loss, name=self.name)

            # Accuracy etc
            self.y_pred_cls = tf.math.argmax(input=self.y_pred, axis=1)
            correct_prediction = tf.math.equal(self.y_pred_cls, self.y_cls)
            self.accuracy = tf.reduce_mean(tf.cast(correct_prediction, dtype=tf.float32)) * 100.

        self._tensorboard()
        tf_utils.show_all_variables(logger=self.logger if self.is_train else None)