Tensorflow 2.0 Keras Модель подклассов

Question

Я пытаюсь реализовать простую UNet-подобную модель, используя метод подклассов модели.Вот мой код:

import tensorflow as tf 
from tensorflow import keras as K

class Enc_block(K.layers.Layer):
    def __init__(self, in_dim):
        super(Enc_block, self).__init__()
        self.conv_layer =  K.layers.SeparableConv2D(in_dim,3, padding='same', activation='relu')
        self.batchnorm_layer = K.layers.BatchNormalization()
        self.pool_layer = K.layers.SeparableConv2D(in_dim,3, padding='same',strides=2, activation='relu')

    def call(self, x):
        x = self.conv_layer(x)
        x = self.batchnorm_layer(x)
        x = self.conv_layer(x)
        x = self.batchnorm_layer(x)
        return self.pool_layer(x), x


class Dec_block(K.layers.Layer):
    def __init__(self, in_dim):
        super(Dec_block, self).__init__()
        self.conv_layer =  K.layers.SeparableConv2D(in_dim,3, padding='same', activation='relu')
        self.batchnorm_layer = K.layers.BatchNormalization()

    def call(self, x):
        x = self.conv_layer(x)
        x = self.batchnorm_layer(x)
        x = self.conv_layer(x)
        x = self.batchnorm_layer(x)
        return x

class Bottleneck(K.layers.Layer):
    def __init__(self, in_dim):
        super(Bottleneck, self).__init__()
        self.conv_1layer =  K.layers.SeparableConv2D(in_dim,1, padding='same', activation='relu')
        self.conv_3layer =  K.layers.SeparableConv2D(in_dim,3, padding='same', activation='relu')
        self.batchnorm_layer = K.layers.BatchNormalization()

    def call(self, x):
        x = self.conv_1layer(x)
        x = self.batchnorm_layer(x)
        x = self.conv_3layer(x)
        x = self.batchnorm_layer(x)
        return x

class Output_block(K.layers.Layer):
    def __init__(self, in_dim):
        super(Output_block, self).__init__()
        self.logits = K.layers.SeparableConv2D(in_dim,3, padding='same', activation=None)
        self.out = K.layers.Softmax()

    def call(self, x):
        x_logits = self.logits(x)
        x = self.out(x_logits)
        return x_logits, x

class UNetModel(K.Model):
    def __init__(self,in_dim):
        super(UNetModel, self).__init__()
        self.encoder_block = Enc_block(in_dim)
        self.bottleneck = Bottleneck(in_dim)
        self.decoder_block = Dec_block(in_dim)
        self.output_block = Output_block(in_dim)


    def call(self, inputs, training=None):
        x, x_skip1 = self.encoder_block(32)(inputs)
        x, x_skip2 = self.encoder_block(64)(x)
        x, x_skip3 = self.encoder_block(128)(x)
        x, x_skip4 = self.encoder_block(256)(x)
        x = self.bottleneck(x)
        x = K.layers.UpSampling2D(size=(2,2))(x)
        x = K.layers.concatenate([x,x_skip4],axis=-1)
        x = self.decoder_block(256)(x)
        x = K.layers.UpSampling2D(size=(2,2))(x) #56x56
        x = K.layers.concatenate([x,x_skip3],axis=-1)
        x = self.decoder_block(128)(x)
        x = K.layers.UpSampling2D(size=(2,2))(x) #112x112
        x = K.layers.concatenate([x,x_skip2],axis=-1)
        x = self.decoder_block(64)(x)
        x = K.layers.UpSampling2D(size=(2,2))(x) #224x224
        x = K.layers.concatenate([x,x_skip1],axis=-1)
        x = self.decoder_block(32)(x)
        x_logits, x = self.output_block(2)(x)
        return x_logits, x

Я получаю следующую ошибку:

ValueError: Input 0 of layer separable_conv2d is incompatible with the layer: expected ndim=4, found ndim=0. Full shape received: []

Я не уверен, что это правильный способ реализации сети в tf.keras. Идеязаключался в реализации блоков кодировщика и декодера путем создания подклассов слоев keras и последующего создания подклассов модели.

Answer 1

Взгляните на эту строку из UNetModel класса:

x, x_skip1 = self.encoder_block(32)(inputs)

, где self.encoder_block() определяется как

self.encoder_block = Enc_block(in_dim)

encoder_block является экземпляром класса. Делая self.encoder_block(32), вы вызываете метод __call__() класса End_block, который ожидает получения повторяемых входных данных изображения rank=4. Вместо этого вы передаете целое число 32, равное rank=0, и вы получаете ValueError, которое говорит именно то, что я только что объяснил: expected ndim=4, found ndim=0. Вероятно, вы намеревались сделать:

x, x_skip1 = self.encoder_block(inputs)

Вы повторяете ту же ошибку и в последующих строках. Существуют дополнительные ошибки, когда вы определяете один и тот же in_dim для каждого пользовательского слоя:

self.encoder_block = Enc_block(in_dim)
self.bottleneck = Bottleneck(in_dim)
self.decoder_block = Dec_block(in_dim)
self.output_block = Output_block(in_dim)

Форма ввода для слоя Bottleneck должна совпадать с формой вывода слоя Enc_Block и т. Д. Я предлагаю вам сначала понять простой пример, прежде чем пытаться реализовать более сложные. Посмотрите на этот пример. Имеет два пользовательских слоя:

import tensorflow as tf
import numpy as np
from tensorflow.keras import layers

class CustomLayer1(layers.Layer):
    def __init__(self, outshape=4):
        super(CustomLayer1, self).__init__()
        self.outshape = outshape
    def build(self, input_shape):
        self.kernel = self.add_weight(name='kernel',
                                      shape=(int(input_shape[1]), self.outshape),
                                      trainable=True)
        super(CustomLayer1, self).build(input_shape)

    def call(self, inputs):
        return tf.matmul(inputs, self.kernel)

class CustomLayer2(layers.Layer):
    def __init__(self):
        super(CustomLayer2, self).__init__()

    def call(self, inputs):
        return inputs / tf.reshape(tf.reduce_sum(inputs, 1), (-1, 1))

Теперь я буду использовать оба этих слоя в новом CombinedLayers классе:

class CombinedLayers(layers.Layer):
    def __init__(self, units=3):
        super(CombinedLayers, self).__init__()
        # `units` defines a number of units in the layer. It is the
        # output shape of the `CustomLayer`
        self.layer1 = CustomLayer1(units) 
        # The input shape is inferred dynamically in the `build()`
        # method of the `CustomLayer1` class
        self.layer2 = CustomLayer1(units)
        # Some layers such as this one do not need to know the shape
        self.layer3 = CustomLayer2()

    def call(self, inputs):
        x = self.layer1(inputs)
        x = self.layer2(x)
        x = self.layer3(x)
        return x

Обратите внимание, что форма ввода CustomLayer1 выводится динамически в методе build(). Теперь давайте проверим это с помощью некоторого ввода:

x_train = [np.random.normal(size=(3, )) for _ in range(5)]
x_train_tensor = tf.convert_to_tensor(x_train)

combined = CombinedLayers(3)

result = combined(x_train_tensor)
result.numpy()
# array([[  0.50822063,  -0.0800476 ,   0.57182697],
#        [ -0.76052217,   0.50127872,   1.25924345],
#        [-19.5887986 ,   9.23529798,  11.35350062],
#        [ -0.33696137,   0.22741248,   1.10954888],
#        [  0.53079047,  -0.08941536,   0.55862488]])

Вот как вы должны подойти к этому. Создайте слои один за другим. Каждый раз, когда вы добавляете новый слой, проверяйте все с некоторыми данными, чтобы убедиться, что вы все делаете правильно.