tf.Module 을 사용하면 tf.Variable 관리를 쉽게 해준다. Model 내 class 변수 Model.trainables, Model.non_trainables, Model.variables 로 모델 내 tf.Variable 을 관리 가능하다.

사용 패턴 1 : static input

class Dense(tf.Module):
    def __init__(self, in_features, out_features, name=None):
        super().__init__(name=name)
        self.w = tf.Variable(
            tf.random.normal([in_features, out_features]), name='w')
        self.b = tf.Variable(tf.zeros([out_features]), name='b')

    def __call__(self, x):
        y = tf.matmul(x, self.w) + self.b
        return tf.nn.relu(y)


class SimpleDNN(tf.Module):
    def __init__(self, name):
        super().__init__(name=name)

        self.dense_1 = Dense(in_features=3, out_features=3)
        self.dense_2 = Dense(in_features=3, out_features=2)

    def __call__(self, x):
        x = self.dense_1(x)
        return self.dense_2(x)

SimpleDNN 안에 2개의 DNN class 에서 생성된 tf.Variable 들을 SimpleDNN instance 내 variables 변수에서 관리 합니다.

그 이외에도 submodules, trainable_variable , non_trainable_variable attribute 가 있습니다.

if __name__ == '__main__':
    simple_dnn = SimpleDNN(name='simple_dnn')
    print('Submodule')
    [print(submodule) for submodule in simple_dnn.submodules]

    print('Variable')
    [print(variable) for variable in simple_dnn.variables]

    print('Trainable variable')
    [print(trainable_variable) for trainable_variable in simple_dnn.trainable_variables]

    print('Non trainable variable')
    [print(non_trainable_variable) for non_trainable_variable in simple_dnn.non_trainable_variables]

사용 패턴 2 : dynamic input

입력 shape 가 동적인 상황에서도 위 모듈을 사용할 수 있습니다.

class FlexibleDenseModule(tf.Module):
    def __init__(self, out_features, name=None):
        super().__init__(name=name)
        self.is_built = False
        self.out_features = out_features

    def __call__(self, x):
        if not self.is_built:
            self.w = tf.Variable(
                tf.random.normal([x.shape[-1], self.out_features]), name='w')
            self.b = tf.Variable(tf.zeros([self.out_features]), name='b')
            self.is_built = True

        y = tf.matmul(x, self.w) + self.b
        return tf.nn.relu(y)

Reference

[tf.Module guide (official)](