tensorflow利用for循环进行训练遇到的内存爆炸问题(OOM)

发布时间：2024-11-29 08:01

文章目录

- - 问题背景
  - 出现问题
  - 解决问题
  - - 部分一
    - 部分二
    - 其他部分
  - 参考

问题背景

最近在用tensorflow学习模型的知识蒸馏，自己基于cifar10数据集训练得到的teacher模型，在对3种不同参数量的student模型使用相同的alpha和temperature参数进行蒸馏之后，得到的实验结果均与论文结果相反（论文：Distilling the Knowledge in a Neural Network）
所以自己打算用for循环方式遍历多种alpha，temperature的参数组合来对比蒸馏效果，然后用matplotlib.pyplot将训练以及对比结果进行绘图（不想自己手动调参了.jpg，在notebook里一遍遍调完参重新跑然后保存数据真难顶）
然后就遇到了很多bug，有tensorflow的out of memory的问题，也有Process和Thread的问题，还有matplotlib.pyplot的问题，没想到代码里面都遇到了

本地环境：
os: win10
显卡：GTX 1650 4G独显
python : 3.8.3
tensorflow : 2.7.0

最后使用的服务器的GPU配置：4个12G的2080Ti（实际只用了GPU：0进行测试）

以下是记录的过程及解决方法

需要用到下面这些模块、函数、类
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
import matplotlib.pyplot as plt
from multiprocessing import Process
from threading import Thread
import time

# 重新Model类
class Distiller(keras.Model):
    def __init__(self, student: keras.Sequential, teacher: keras.Sequential):
        super().__init__()
        self.teacher = teacher
        self.student = student

    def compile(self,
                optimizer: keras.optimizers.Optimizer,
                metrics,
                student_loss_fn,
                distillation_loss_fn,
                alpha=0.1,
                temperature=3):
        ......

    def train_step(self, data):
        x, y = data
        # teacher prediction
        ......
        return results

    def test_step(self, data):
        x, y = data
        ......
        return results

# 创建模型
def build_model(name, conv1_size, conv2_size, conv3_size, dense_size):
    model = keras.Sequential([
        ......
    ],
        name=name)

    return model

循环使用的主函数
def main_loop(alpha, T):
    # 将loop_param也设置为全局变量，不然无法return此参数
    global loop_param
    loop_param = str(alpha) + \'_\' + str(T) + \'_\'
    # 打印本次循环的信息
    print(\'\\n\\n\' + loop_param)

    # 创建student
    student = build_model(\'student\', 8, 16, 16, 16)
    # 记录student初始结构
    student_scratch = keras.models.clone_model(student)
    ......

def draw_distill():
......
def draw_scratch():
......

出现问题

一开始直接在for循环中进行distiller和student_scratch的训练与评估，之后几个循环之后就出现了OOM，部分输出如下：

......
Epoch 2/20
2022-03-23 21:34:09.303394: W tensorflow/core/common_runtime/bfc_allocator.cc:462] Allocator (GPU_0_bfc) ran out of memory trying to allocate 242.0KiB (rounded to 247808)requested by op student/conv2d_2/Relu
If the cause is memory fragmentation maybe the environment variable \'TF_GPU_ALLOCATOR=cuda_malloc_async\' will improve the situation.
Current allocation summary follows.
......
2022-03-23 21:34:09.332826: I tensorflow/core/common_runtime/bfc_allocator.cc:1078] Sum Total of in-use chunks: 2.10GiB
2022-03-23 21:34:09.332899: I tensorflow/core/common_runtime/bfc_allocator.cc:1080] total_region_allocated_bytes_: 2255958784 memory_limit_: 2255958836 available bytes: 52 curr_region_allocation_bytes_: 4511918080
2022-03-23 21:34:09.332977: I tensorflow/core/common_runtime/bfc_allocator.cc:1086] Stats:
Limit:                      2255958836
InUse:                      2255956224
MaxInUse:                   2255956224
NumAllocs:                    55811090
MaxAllocSize:                614400000
Reserved:                            0
PeakReserved:                        0
LargestFreeBlock:                    0
......
tensorflow.python.framework.errors_impl.InternalError: Failed copying input tensor from /job:localhost/replica:0/task:0/device:CPU:0 to /job:localhost/replica:0/task:0/device:GPU:0 in order to run _EagerConst: Dst tensor is not initialized.

所以想着用Thread或者Process来将每个循环进行隔离，使用单进程或者单线程。这样每次循环结束之后创建的进程 / 线程也会终止，tensorflow占用的空间就不会一直增加直到OOM

解决问题

然后尝试使用Thread或者Process来为每次循环创建单独的进程或线程

部分一

使用Thread，然后模型的训练、评估以及用plt将训练结果绘图等都在main_loop函数中进行，但遇到了：
UserWarning: Starting a Matplotlib GUI outside of the main thread will likely fail的Warning。不过虽然提示说will likely fail，但自己实验中用plt还是能够正常绘图与保存图片的。
以防万一，如果要避免此Warning的提示，可以参照另一篇文章：pyplot.plot使用遇到：UserWarning: Starting a Matplotlib GUI outside of the main thread will likely fail

if __name__ == \'__main__\':
    # 加载数据集
    (train_images, train_labels), (test_images, test_labels) = keras.datasets.cifar10.load_data()
    # Normalize pixel values
    train_images, test_images = train_images / 255.0, test_images / 255.0

    teacher = build_model(\'teacher\', 32, 64, 64, 64)
    # load teacher model from SavedModel
    teacher = keras.models.load_model(\'teacher_model\')

    # 利用循环调参，观察不同超参数对应的蒸馏效果
    for alpha in (0.1, 0.2, 0.3):
        for T in range(5, 21, 5):
            print(time.strftime(\'%H-%M-%S: \'))
            t = Thread(target=main_loop, args=(alpha, T))
            t.start()
            t.join()
            draw_distill()
            draw_scratch()

部分二

使用Process
与上面代码内容基本相同，只是将Thread换成Process，然后把需要的参数都传入Process的args中即可。此时可以直接将模型的训练、评估以及用plt把训练结果进行绘图等都放在main_loop函数中进行。

简单的写一下更改后的代码：
if __name__ == \'__main__\':
    # 加载数据集
    ......
    # 利用循环调参，观察不同超参数对应的蒸馏效果
    for alpha in (0.1, 0.2, 0.3):
        for T in range(5, 21, 5):
            print(time.strftime(\'%H-%M-%S: \'))
            因为创建的子进程与主进程是相互隔离的，所以无法像子线程直接使用主线程的变量数据那样，只能把需要的变量作为参数传入Process中
            p = Process(target=main_loop, args=(alpha, T, teacher, train_images, train_labels, test_images, test_labels))
            p.start()
            # 阻塞当前进程，告诉主进程，当进程p执行完之后才能向后继续执行进程
            p.join()

更详细的Thread与Process的区别以及简单的使用方法写在了另一篇文章：
python multiprocessing.Process与threading.Thread的区别以及多进程，多线程的一些使用方法

其他部分

此外，在本地机器上尝试了用GPU来进行训练，但是很慢，而且输出很多过程的信息。输出的部分信息如下：

 376/1563 [======>.......................] - ETA: 34s - loss: 2.1089 - accuracy: 0.20432022-03-23 11:18:02.870737: I tensorflow/core/common_runtime/eager/execute.cc:1224] Executing op ReadVariableOp in device /job:localhost/replica:0/task:0/device:GPU:0
2022-03-23 11:18:02.871020: I tensorflow/core/common_runtime/eager/execute.cc:1224] Executing op Identity in device /job:localhost/replica:0/task:0/device:GPU:0
......
2022-03-23 11:18:02.871797: I 
tensorflow/core/common_runtime/eager/execute.cc:1224] Executing op Identity in device /job:localhost/replica:0/task:0/device:GPU:0
2022-03-23 11:18:02.925364: I tensorflow/core/common_runtime/eager/execute.cc:1224] Executing op __inference_train_function_914 in device /job:localhost/replica:0/task:0/device:GPU:0
 383/1563 [======>.......................] - ETA: 34s - loss: 2.1056 - accuracy: 0.20602022-03-23 11:18:02.931965: I tensorflow/core/common_runtime/eager/execute.cc:1224] Executing op ReadVariableOp in device /job:localhost/replica:0/task:0/device:GPU:0
......
2022-03-23 11:18:02.971101: I tensorflow/core/common_runtime/eager/execute.cc:1224] Executing op __inference_train_function_914 in device /job:localhost/replica:0/task:0/device:GPU:0
 391/1563 [======>.......................] - ETA: 33s - loss: 2.1030 - accuracy: 0.20772022-03-23 11:18:02.977031: I tensorflow/core/common_runtime/eager/execute.cc:1224] Executing op ReadVariableOp in device /job:localhost/replica:0/task:0/device:GPU:0
......
2022-03-23 11:18:03.020212: I tensorflow/core/common_runtime/eager/execute.cc:1224] Executing op __inference_train_function_914 in device /job:localhost/replica:0/task:0/device:GPU:0
 400/1563 [======>.......................] - ETA: 32s - loss: 2.0986 - accuracy: 0.20922022-03-23 11:18:03.025223: I tensorflow/core/common_runtime/eager/execute.cc:1224] Executing op ReadVariableOp in device /job:localhost/replica:0/task:0/device:GPU:0
 ......

即将写完本文之时，又参考别的blog(参考部分有链接)里面说的方法用GPU训练了一下，在import tensorflow as tf之前，设置：
os.environ[\'TF_CPP_MIN_LOG_LEVEL\'] = \'3\'，之后发现果然有效，很多无用的信息都不再输出。
还尝试在服务器上用GPU进行训练（之前没用服务器测试，就是因为上述原因，输出了太多冗余信息），比本地的GPU快了很多：

服务器：GPU total train time: 62.842313051223755 s
本地：GPU total train time: 104.92916321754456 s

参考

https://keras.io/examples/vision/knowledge_distillation
https://www.tensorflow.org/guide/gpu?hl=zh-cn
https://blog.csdn.net/dcrmg/article/details/80029741