请提供一些关于您想法的最小评论，以便我改进我的查询。谢谢。-)
我正在尝试使用梯度累积（GA）训练一个tf.keras模型。但我不想在自定义训练循环中使用它（像这样），而是通过重写train_step来自定义.fit()方法。这可能吗？如何实现？原因是如果我们想利用keras内置功能，如fit、callbacks，我们不想使用自定义训练循环，但与此同时，如果我们出于某些原因（如GA或其他）想重写train_step，我们可以自定义fit方法，同时仍然利用这些内置功能的优势。
此外，我知道使用GA的优点，但使用它的主要缺点是什么？为什么它不是框架的默认功能，而是一个可选功能？
# 重写训练步骤 # 我的尝试 # 它没有适当实现 # 需要修复 class CustomTrainStep(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) # 梯度带 with tf.GradientTape() as tape: y_pred = self(x, training=True) loss = self.compiled_loss( y, y_pred, regularization_losses=self.losses ) # 计算批次梯度 gradients = tape.gradient(loss, self.trainable_variables) # 累积批次梯度 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 ] # 应用累积梯度 self.optimizer.apply_gradients( zip(accum_gradient, self.trainable_variables) ) # TODO: 重置 self.gradient_accumulation # 更新指标 self.compiled_metrics.update_state(y, y_pred) return {m.name: m.result() for m in self.metrics}请运行并使用以下玩具设置进行检查。
# 模型大小 = 32input = keras.Input(shape=(size,size,3))efnet = keras.applications.DenseNet121( weights=None, include_top = False, input_tensor = input)base_maps = keras.layers.GlobalAveragePooling2D()(efnet.output) base_maps = keras.layers.Dense( units=10, activation='softmax', name='primary')(base_maps)custom_model = CustomTrainStep( n_gradients=10, inputs=[input], outputs=[base_maps])# 绑定所有custom_model.compile( loss = keras.losses.CategoricalCrossentropy(), metrics = ['accuracy'], optimizer = keras.optimizers.Adam())# 数据 (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]) # 如果我们想调整大小 y_train = tf.one_hot(y_train , depth=10) # 自定义 fit custom_model.fit(x_train, y_train, batch_size=64, epochs=3, verbose = 1)更新
我发现其他人也尝试实现这一点，并且遇到了相同的问题。有人找到了一些解决方法，在这里，但它太混乱了，我认为应该有更好的方法。
更新 2
接受的答案（由Mr.For Example提供）很好，在单一策略中工作良好。现在，我喜欢开始第二个悬赏，以扩展它以支持多GPU、TPU以及混合精度技术。有一些复杂情况，参见详情。
回答：
是的，通过重写train_step，可以在不使用自定义训练循环的情况下自定义.fit()方法，以下简单示例将向您展示如何使用梯度累积训练一个简单的MNIST分类器：
import tensorflow as tf class CustomTrainStep(tf.keras.Model): def __init__(self, n_gradients, *args, **kwargs): super().__init__(*args, **kwargs) self.n_gradients = tf.constant(n_gradients, dtype=tf.int32) self.n_acum_step = tf.Variable(0, dtype=tf.int32, trainable=False) self.gradient_accumulation = [tf.Variable(tf.zeros_like(v, dtype=tf.float32), trainable=False) for v in self.trainable_variables] def train_step(self, data): self.n_acum_step.assign_add(1) x, y = data # 梯度带 with tf.GradientTape() as tape: y_pred = self(x, training=True) loss = self.compiled_loss(y, y_pred, regularization_losses=self.losses) # 计算批次梯度 gradients = tape.gradient(loss, self.trainable_variables) # 累积批次梯度 for i in range(len(self.gradient_accumulation)): self.gradient_accumulation[i].assign_add(gradients[i]) # 如果 n_acum_step 达到 n_gradients，则我们应用累积梯度来更新变量，否则什么也不做 tf.cond(tf.equal(self.n_acum_step, self.n_gradients), self.apply_accu_gradients, lambda: None) # 更新指标 self.compiled_metrics.update_state(y, y_pred) return {m.name: m.result() for m in self.metrics} def apply_accu_gradients(self): # 应用累积梯度 self.optimizer.apply_gradients(zip(self.gradient_accumulation, self.trainable_variables)) # 重置 self.n_acum_step.assign(0) for i in range(len(self.gradient_accumulation)): self.gradient_accumulation[i].assign(tf.zeros_like(self.trainable_variables[i], dtype=tf.float32))# 模型 input = tf.keras.Input(shape=(28, 28))base_maps = tf.keras.layers.Flatten(input_shape=(28, 28))(input)base_maps = tf.keras.layers.Dense(128, activation='relu')(base_maps)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])# 绑定所有custom_model.compile( loss = tf.keras.losses.CategoricalCrossentropy(), metrics = ['accuracy'], optimizer = tf.keras.optimizers.Adam(learning_rate=1e-3) )# 数据 (x_train, y_train), (_, _) = tf.keras.datasets.mnist.load_data()x_train = tf.divide(x_train, 255)y_train = tf.one_hot(y_train , depth=10) # 自定义 fit custom_model.fit(x_train, y_train, batch_size=6, epochs=3, verbose = 1)