使用tf.GradientTape()训练模型，模型却不会收敛，loss一直大幅度波动或者不变得一些解决办法

一、提出问题

import tensorflow.keras as keras
import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
import nets_VAE

if __name__ == '__main__':
    batch_size = 2
    epoch = 100
    iteration = 30000
    (train_x, train_y), (test_x, test_y) = keras.datasets.mnist.load_data()
    train_x = (train_x / 255.0).reshape(-1, 28, 28, 1)
    test_x = (test_x / 255.0).reshape(-1, 28, 28, 1)
    model = nets_VAE.Classifier().model()
    opt = keras.optimizers.Adam(lr=0.002)
    SCC = keras.losses.SparseCategoricalCrossentropy()
    model.compile(optimizer=opt, loss=SCC, metrics=['accuracy'])
    # 执行gradientTape()方案
    for i in range(100):
        print('epoch=' + str(i))
        for j in range(iteration):
            with tf.GradientTape() as g:
                start = j * batch_size
                end = start + batch_size
                img_x = train_x[start:end]
                img_y = train_y[start:end]
                y = model(img_x)
                loss = SCC(img_y, y)
                grads = g.gradient(loss, model.trainable_variables)
                opt.apply_gradients(zip(grads, model.trainable_variables))
                if j % 200 == 0:
                    model.evaluate(x=test_x[:1000], y=test_y[:1000])
    # 执行tensorflow内部训练方案
    # model.fit(x=train_x, y=train_y, epochs=100, validation_data=(test_x, test_y))

直接执行以上代码你大概率会得到以下结果：

可以发现，在训练一切正常的情况下，模型的准确率几乎不会变，并且还非常低。我曾经考虑过几种情况：

1、会不会网络结构设计不合理？

2、是不是使用梯度带（GradientTape）的方式有错？

3、代码写错了？

我首先使用Model().fit()方法对模型训练，一切正常，所以排除第一种可能。然后对比以前写的正常的代码及认真核对后，后两者也被排除。那么是什么原因导致在一切都看似正常的情况下，使用GradientTape()却得不到正常的结果呢？

二、找原因

经过再一次分析，我使用GradientTape()对模型进行训练时，batch-size的大小是我自定义的，而Model().fit()是自动分配的（本例中为32）。因此我猜测原因可能是batch-size大小导致的。于是相应的我对其进行了修改。

batch_size = 50
epoch = 100
iteration = 1200

本次训练得到的结果如下：

可以发现一切都变得可以接受了，但是为什么呢？batch-size为什么会对训练有这么大的影响呢？具体可以看看这两篇文章 https://blog.csdn.net/stay_foolish12/article/details/91984832和 https://blog.csdn.net/qq_40636392/article/details/98750938

batch-size过小时，样本集可能无法代表全部样本的统计规律，得到的梯度可能各自为政，相互抵消。因此对模型训练起到反作用。

batch-size过大，一般来说对模型训练不会起到太大的负面作用，有时甚至会比分批训练好，但是它可能会在一定程度上增大训练周期。

其次，batch-size的大小会影响到梯度的大小。当batch-size过小时，学习率往往不宜过大，否则会导致梯度波动过大。在上面的例子中，倘若我们不改动batch-size的大小为50，而将学习率调整为0.0002也可以获得正确结果。

三、解决方案

综上所述，当使用GradientTape()训练模型时,应该注意以下几点：

1、batch-size不易过小，条件允许的情况下。50左右差不多。

2、batch-size已经达到电脑内存极限，但训练效果仍然很差时应该适当调低学习率。