在 Tensorflow 中实现 UNet

Implement UNet in Tensorflow

我正在尝试使用 Keras API 在 TensorFlow 2 中实现 UNet 图像分割,但我不确定如何实现 Concatenate 层。这是我尝试过的:

def create_model_myunet(depth, start_f, output_channels, encoder_kernel_size):

    # Encoder
    model = tf.keras.Sequential()

    for i in range(0, depth):

        if i == 0:
            print("Specifying an input shape")
            input_shape = [config.img_h, config.img_w, 3]
        else:
            input_shape = [None]

        model.add(tf.keras.layers.Conv2D(filters=2**(start_f+i), 
                                         kernel_size=(encoder_kernel_size, encoder_kernel_size),
                                         strides=(1, 1),
                                         padding='same',
                                         input_shape=input_shape,
                                         name = "enc_conv2d_" + str(i)))
        model.add(tf.keras.layers.ReLU(name = "enc_relu_" + str(i)))
        model.add(tf.keras.layers.MaxPool2D(pool_size=(2, 2), name="enc_maxpool2d_" + str(i)))

    # Decoder
    initializer = tf.random_normal_initializer(0., 0.02)

    for i in range(depth, 1, -1):

        model.add(
        tf.keras.layers.Conv2DTranspose(2**(start_f+i),
                                        encoder_kernel_size,
                                        strides=2,
                                        padding='same',
                                        kernel_initializer=initializer,
                                        use_bias=False)
        )

        model.add(tf.keras.layers.BatchNormalization())

        model.add(tf.keras.layers.ReLU(name="dec_relu_"+str(i)))

        model.add(tf.keras.layers.Concatenate([
            model.get_layer(name="dec_relu_"+str(i)).output,  
            model.get_layer(name="enc_relu_"+str(i-1)).output
        ] ))
        pass

    last = tf.keras.layers.Conv2DTranspose(
      output_channels, 3, strides=2,
      padding='same', activation='softmax')  #64x64 -> 128x128

    model.add(last)

    return model

它给我以下错误:

ValueError: A Concatenate layer should be called on a list of at least 2 inputs

你需要改变

model.add(tf.keras.layers.Concatenate([
    model.get_layer(name="dec_relu_"+str(i)).output,  
    model.get_layer(name="enc_relu_"+str(i-1)).output
] ))

model.add(tf.keras.layers.Concatenate()([  # Sequential api
    model.get_layer(name="dec_relu_"+str(i)).output,  
    model.get_layer(name="enc_relu_"+str(i-1)).output
] ))

model.add(tf.keras.layers.concatenate([  # Functional api
    model.get_layer(name="dec_relu_"+str(i)).output,  
    model.get_layer(name="enc_relu_"+str(i-1)).output
] ))