您好捧毛，我是碼農(nóng)飛哥，感謝您閱讀本文让网，歡迎一鍵三連哦呀忧。
本文只是記錄我優(yōu)化的心酸歷程。無(wú)他溃睹，唯記錄爾而账。。因篇。泞辐。。小伙伴們可圍觀竞滓，可打call咐吼，可以私信與我交流。
干貨滿滿商佑，建議收藏锯茄，需要用到時(shí)常看看茶没。 小伙伴們?nèi)缬袉?wèn)題及需要肌幽，歡迎踴躍留言哦~ ~ ~。

問(wèn)題背景

現(xiàn)有一個(gè)古詩(shī)自動(dòng)生成的訓(xùn)練接口礁叔，該接口通過(guò)Pytorch來(lái)生訓(xùn)練模型（即生成古詩(shī)）為了加速使用到了GPU牍颈，但是訓(xùn)練完成之后GPU未能釋放。故此需要進(jìn)行優(yōu)化琅关，即在古詩(shī)生成完成之后釋放GPU煮岁。
該項(xiàng)目是一個(gè)通過(guò)Flask搭建的web服務(wù)讥蔽，在服務(wù)器上為了實(shí)現(xiàn)并發(fā)采用的是gunicorn來(lái)啟動(dòng)應(yīng)用。通過(guò)pythorch來(lái)進(jìn)行古詩(shī)訓(xùn)練画机。項(xiàng)目部署在一個(gè)CentOS的服務(wù)器上冶伞。

系統(tǒng)環(huán)境

因?yàn)樘厥獾脑蜻@里之后一個(gè)服務(wù)器供使用步氏，故不能考慮加機(jī)器的情況响禽。

優(yōu)化歷程

pytorch在訓(xùn)練模型時(shí)，需要先加載模型model和數(shù)據(jù)data荚醒，如果有GPU顯存的話我們可以將其放到GPU顯存中加速芋类，如果沒(méi)有GPU的話則只能使用CPU了。
由于加載模型以及數(shù)據(jù)的過(guò)程比較慢界阁。所以侯繁，我這邊將加載過(guò)程放在了項(xiàng)目啟動(dòng)時(shí)加載。

import torch
device = "cuda" if torch.cuda.is_available() else "cpu"
model = GPT2LMHeadModel.from_pretrained(os.path.join(base_path, "model"))
model.to(device)
model.eval()

這部分耗時(shí)大約在6秒左右泡躯。cuda表示使用torch的cuda贮竟。模型數(shù)據(jù)加載之后所占的GPU顯存大小大約在1370MB。優(yōu)化的目標(biāo)就是在訓(xùn)練完成之后將這部分占用的顯存釋放掉较剃。

小小分析一波

現(xiàn)狀是項(xiàng)目啟動(dòng)時(shí)就加載模型model和數(shù)據(jù)data的話咕别，當(dāng)模型數(shù)據(jù)在GPU中釋放掉之后，下次再進(jìn)行模型訓(xùn)練的話不就沒(méi)有模型model和數(shù)據(jù)data了么写穴？如果要釋放GPU的話惰拱，就需要考慮如何重新加載GPU。
所以啊送，模型model和數(shù)據(jù)data不能放在項(xiàng)目啟動(dòng)的時(shí)候加載弓颈，只能放在調(diào)用訓(xùn)練的函數(shù)時(shí)加載，但是由于加載比較慢删掀，所以只能放在一個(gè)異步的子線程或者子進(jìn)程中運(yùn)行翔冀。
所以，我這邊首先將模型數(shù)據(jù)的加載過(guò)程以及訓(xùn)練放在了一個(gè)單獨(dú)的線程中執(zhí)行披泪。

第一階段：直接上torch.cuda.empty_cache()清理纤子。

GPU沒(méi)釋放，那就釋放唄款票。這不是很簡(jiǎn)單么控硼？百度一波pytorch怎么釋放GPU顯存。

在這里插入圖片描述

在這里插入圖片描述

第二階段（創(chuàng)建子進(jìn)程加載模型并進(jìn)行訓(xùn)練）

既然子線程加載模型并進(jìn)行訓(xùn)練不能釋放GPU的話悉抵，那么我們能不能轉(zhuǎn)變一下思路肩狂。創(chuàng)建一個(gè)子進(jìn)程來(lái)加載模型數(shù)據(jù)并進(jìn)行訓(xùn)練，
當(dāng)訓(xùn)練完成之后就將這個(gè)子進(jìn)程殺掉姥饰，它所占用的資源（主要是GPU顯存）不就被釋放了么婚温？
這思路看起來(lái)沒(méi)有絲毫的毛病呀。說(shuō)干就干媳否。

1. 定義加載模型數(shù)據(jù)以及訓(xùn)練的方法 training。（代碼僅供參考）

def training(queue):
    manage.app.app_context().push()
    current_app.logger.error('基礎(chǔ)加載開(kāi)始')
    with manage.app.app_context():
        device = "cuda" if torch.cuda.is_available() else "cpu"
        current_app.logger.error('device1111開(kāi)始啦啦啦')
        model.to(device)
        current_app.logger.error('device2222')
        model.eval()
        n_ctx = model.config.n_ctx
        current_app.logger.error('基礎(chǔ)加載完成')
        #訓(xùn)練方法
        result_list=start_train(model,n_ctx,device)
        current_app.logger.error('完成訓(xùn)練')
        #將訓(xùn)練方法返回的結(jié)果放入隊(duì)列中
        queue.put(result_list)      

2. 創(chuàng)建子進(jìn)程執(zhí)行training方法荆秦，然后通過(guò)阻塞的方法獲取訓(xùn)練結(jié)果

from torch import multiprocessing as mp
def sub_process_train():
    #定義一個(gè)隊(duì)列獲取訓(xùn)練結(jié)果
    train_queue = mp.Queue()
    training_process = mp.Process(target=training, args=(train_queue))
    training_process.start()
    current_app.logger.error('子進(jìn)程執(zhí)行')
    # 等訓(xùn)練完成
    training_process.join()
    current_app.logger.error('執(zhí)行完成')
    #獲取訓(xùn)練結(jié)果
    result_list = train_queue.get()
    current_app.logger.error('獲取到數(shù)據(jù)')
    if training_process.is_alive():
        current_app.logger.error('子進(jìn)程還存活')
        #殺掉子進(jìn)程
        os.kill(training_process.pid, signal.SIGKILL)
    current_app.logger.error('殺掉子進(jìn)程')
    return result_list
        

3. 因?yàn)樽舆M(jìn)程要阻塞獲取執(zhí)行結(jié)果篱竭，所以需要定義一個(gè)線程去執(zhí)行sub_process_train方法以保證訓(xùn)練接口可以正常返回。

import threading
threading.Thread(target=sub_process_train).start()

代碼寫好了步绸，見(jiàn)證奇跡的時(shí)候來(lái)了掺逼。

首先用python manage.py 啟動(dòng)一下，看下結(jié)果瓤介，運(yùn)行結(jié)果如下吕喘，報(bào)了一個(gè)錯(cuò)誤，從錯(cuò)誤的提示來(lái)看就是不能在forked的子進(jìn)程中重復(fù)加載CUDA刑桑。"Cannot re-initialize CUDA in forked subprocess. " + msg) RuntimeError: Cannot re-initialize CUDA in forked subprocess. To use CUDA with multiprocessing, you must use the 'spawn' start method氯质。

在這里插入圖片描述

fork和spawn是構(gòu)建子進(jìn)程的不同方式琢锋，區(qū)別在于
1. fork: 除了必要的啟動(dòng)資源辕漂，其余的變量，包吴超，數(shù)據(jù)等都集成自父進(jìn)程钉嘹，也就是共享了父進(jìn)程的一些內(nèi)存頁(yè)，因此啟動(dòng)較快鲸阻，但是由于大部分都是用的自父進(jìn)程數(shù)據(jù)跋涣，所有是不安全的子進(jìn)程缨睡。
2. spawn：從頭構(gòu)建一個(gè)子進(jìn)程，父進(jìn)程的數(shù)據(jù)拷貝到子進(jìn)程的空間中仆潮，擁有自己的Python解釋器宏蛉，所有需要重新加載一遍父進(jìn)程的包，因此啟動(dòng)叫慢性置，但是由于數(shù)據(jù)都是自己的拾并，安全性比較高。

回到剛剛那個(gè)報(bào)錯(cuò)上面去鹏浅。為啥提示要不能重復(fù)加載嗅义。

這是因?yàn)镻ython3中使用 spawn啟動(dòng)方法才支持在進(jìn)程之間共享CUDA張量。而用的multiprocessing 是使用 fork 創(chuàng)建子進(jìn)程隐砸，不被 CUDA 運(yùn)行時(shí)所支持之碗。
所以，只有在創(chuàng)建子進(jìn)程之前加上mp.set_start_method('spawn') 方法季希。即

def sub_process_train(prefix, length):
    try:
        mp.set_start_method('spawn')
    except RuntimeError:
        pass
    train_queue = mp.Queue()
    training_process = mp.Process(target=training, args=(train_queue))
    ##省略其他代碼

再次通過(guò) python manage.py 運(yùn)行項(xiàng)目褪那。運(yùn)行結(jié)果圖1和圖2所示，可以看出可以正確是使用GPU顯存式塌，在訓(xùn)練完成之后也可以釋放GPU博敬。

圖1

圖2

在這里插入圖片描述

第三階段（全局線程池+釋放GPU）

子進(jìn)程的方式也不行了括勺。只能回到前面的線程方式了。前面創(chuàng)建線程的方式都是直接通過(guò)直接new一個(gè)新線程的方式曲掰，當(dāng)同時(shí)運(yùn)行的線程數(shù)過(guò)多的話疾捍，則很容易就會(huì)出現(xiàn)GPU占滿的情況，從而導(dǎo)致應(yīng)用崩潰栏妖。所以乱豆，這里采用全局線程池的方式來(lái)創(chuàng)建并管理線程，然后當(dāng)線程執(zhí)行完成之后釋放資源吊趾。

1. 在項(xiàng)目啟動(dòng)之后就創(chuàng)建一個(gè)全局線程池宛裕。大小是2瑟啃。保證還有剩余的GPU。

from multiprocessing.pool import ThreadPool
pool = ThreadPool(processes=2)

2. 通過(guò)線程池來(lái)執(zhí)行訓(xùn)練

  pool.apply_async(func=async_produce_poets)

3. 用線程加載模型和釋放GPU

def async_produce_poets():
    try:
        print("子進(jìn)程開(kāi)始" + str(os.getpid())+" "+str(threading.current_thread().ident))
        start_time = int(time.time())
        manage.app.app_context().push()
        device = "cuda" if torch.cuda.is_available() else "cpu"
        model = GPT2LMHeadModel.from_pretrained(os.path.join(base_path, "model"))
        model.to(device)
        model.eval()
        n_ctx = model.config.n_ctx
        result_list=start_train(model,n_ctx,device)
        #將模型model轉(zhuǎn)到cpu
        model = model.to('cpu')
        #刪除模型揩尸，也就是刪除引用
        del model
        #在使用其釋放GPU蛹屿。
        torch.cuda.empty_cache()
        train_seconds = int(time.time() - start_time)
        current_app.logger.info('訓(xùn)練總耗時(shí)是={0}'.format(str(train_seconds)))
    except Exception as e:
        manage.app.app_context().push()

這一番操作之后，終于達(dá)到了理想的效果岩榆。

這里因?yàn)槭褂玫搅薵unicorn來(lái)啟動(dòng)項(xiàng)目错负。所以gunicorn 相關(guān)的知識(shí)必不可少。在CPU受限的系統(tǒng)中采用sync的工作模式比較理想勇边。
詳情可以查看gunicorn的簡(jiǎn)單總結(jié)

問(wèn)題分析犹撒，前面第一階段直接使用torch.cuda.empty_cache() 沒(méi)能釋放GPU就是因?yàn)闆](méi)有刪除掉模型model。模型已經(jīng)加載到了GPU了粒褒。

總結(jié)

本文從實(shí)際項(xiàng)目的優(yōu)化入手识颊，記錄優(yōu)化方面的方方面面。希望對(duì)讀者朋友們有所幫助奕坟。

參考

multiprocessing fork() vs spawn()

粉絲專屬福利

軟考資料：實(shí)用軟考資料

面試題：5G 的Java高頻面試題

學(xué)習(xí)資料：50G的各類學(xué)習(xí)資料

脫單秘籍：回復(fù)【脫單】

并發(fā)編程：回復(fù)【并發(fā)編程】

全網(wǎng)同名【碼農(nóng)飛哥】祥款。不積跬步，無(wú)以至千里月杉，享受分享的快樂(lè)
我是碼農(nóng)飛哥刃跛，再次感謝您讀完本文。