【发布时间】：2021-06-02 22:23:53
【问题描述】：
请对您的想法添加最低限度的评论，以便我改进查询。谢谢你。 -)
我正在尝试使用 梯度累积 (GA) 训练 tf.keras 模型。但是我不想在自定义训练循环（like）中使用它，而是通过覆盖train_step来自定义.fit()方法。这可能吗？如何做到这一点？原因是如果我们想获得 keras 内置功能（如 fit、callbacks）的好处，我们不想使用自定义训练循环，但同时如果我们想覆盖 @987654329 @ 出于某种原因（例如 GA 或其他），我们可以自定义 fit 方法，并且仍然可以利用这些内置函数。
而且，我知道使用 GA 的优点，但使用它的主要缺点是什么？为什么它不是作为框架的默认功能而是作为可选功能提供的？
# overriding train step
# my attempt
# it's not appropriately implemented
# and need to fix
class CustomTrainStep(tf.keras.Model):
def __init__(self, n_gradients, *args, **kwargs):
super().__init__(*args, **kwargs)
self.n_gradients = n_gradients
self.gradient_accumulation = [tf.zeros_like(this_var) for this_var in \
self.trainable_variables]
def train_step(self, data):
x, y = data
batch_size = tf.cast(tf.shape(x)[0], tf.float32)
# Gradient Tape
with tf.GradientTape() as tape:
y_pred = self(x, training=True)
loss = self.compiled_loss(y, y_pred, regularization_losses=self.losses)
# Calculate batch gradients
gradients = tape.gradient(loss, self.trainable_variables)
# Accumulate batch gradients
accum_gradient = [(acum_grad+grad) for acum_grad, grad in \
zip(self.gradient_accumulation, gradients)]
accum_gradient = [this_grad/batch_size for this_grad in accum_gradient]
# apply accumulated gradients
self.optimizer.apply_gradients(zip(accum_gradient, self.trainable_variables))
# TODO: reset self.gradient_accumulation
# update metrics
self.compiled_metrics.update_state(y, y_pred)
return {m.name: m.result() for m in self.metrics}
请运行并检查以下玩具设置。
# Model
size = 32
input = tf.keras.Input(shape=(size,size,3))
efnet = tf.keras.applications.DenseNet121(weights=None,
include_top = False,
input_tensor = input)
base_maps = tf.keras.layers.GlobalAveragePooling2D()(efnet.output)
base_maps = tf.keras.layers.Dense(units=10, activation='softmax',
name='primary')(base_maps)
custom_model = CustomTrainStep(n_gradients=10, inputs=[input], outputs=[base_maps])
# bind all
custom_model.compile(
loss = tf.keras.losses.CategoricalCrossentropy(),
metrics = ['accuracy'],
optimizer = tf.keras.optimizers.Adam() )
# data
(x_train, y_train), (_, _) = tf.keras.datasets.mnist.load_data()
x_train = tf.expand_dims(x_train, -1)
x_train = tf.repeat(x_train, 3, axis=-1)
x_train = tf.divide(x_train, 255)
x_train = tf.image.resize(x_train, [size,size]) # if we want to resize
y_train = tf.one_hot(y_train , depth=10)
# customized fit
custom_model.fit(x_train, y_train, batch_size=64, epochs=3, verbose = 1)
更新
我发现其他一些人也试图实现这一目标并最终遇到了同样的问题。有一些解决方法，here，但它太乱了，我认为应该有一些更好的方法。
【问题讨论】：
标签： python tensorflow machine-learning keras deep-learning