tf.keras 中training的使用

import tensorflow as f
from tensorflow.keras import layers, losses


class OneLayer(layers.Layer):

    def __init__(self):
        super(OneLayer, self).__init__()

    # def call(self, inputs, **kwargs):
    # def call(self, inputs, **kwargs):
    #
    #     if kwargs.get("training", False):
    #         inputs = inputs + 1
    #         tf.print("加11111111111111111111111111111111")
    #     else:
    #         inputs = inputs - 1
    #         tf.print("减11111111111111111111111111111111")
    #
    #     return inputs
    #

    def call(self, inputs, training=None, mask=None, **kwargs):

        if training:
            inputs = inputs + 1
            tf.print("加11111111111111111111111111111111")
        else:
            inputs = inputs - 1
            tf.print("减11111111111111111111111111111111")

        return inputs



inputs = layers.Input((1, ))

embedding = layers.Dense(64)(inputs)
embedding = OneLayer()(embedding)

embedding = tf.multiply(embedding, 2.0)

output = layers.Dense(1)(embedding)

model = tf.keras.Model(inputs, output)

model.summary()

import numpy as np

x = np.arange(64).reshape((64, 1)) * 1.0
y = np.arange(64).reshape((64, 1)) * 1.0 + 1



model.compile("adam", loss=losses.MeanSquaredError(), metrics=["mse"])

model.fit(x, y, batch_size=16, epochs=100)


print("----------------------------")

print(model.predict(x[:2]))