您好，我是碼農(nóng)飛哥，感謝您閱讀本文，歡迎一鍵三連哦。
本文只是記錄我優(yōu)化的心酸歷程。無(wú)他，唯記錄爾。。。。。小伙伴們可圍觀(guān)，可打call，可以私信與我交流。
干貨滿(mǎn)滿(mǎn)，建議收藏，需要用到時(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顯存的話(huà)我們可以將其放到GPU顯存中加速，如果沒(méi)有GPU的話(huà)則只能使用CPU了。
由于加載模型以及數(shù)據(jù)的過(guò)程比較慢。所以，我這邊將加載過(guò)程放在了項(xiàng)目啟動(dòng)時(shí)加載。

import torchdevice = "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的話(huà)，當(dāng)模型數(shù)據(jù)在GPU中釋放掉之后，下次再進(jìn)行模型訓(xùn)練的話(huà)不就沒(méi)有模型model和數(shù)據(jù)data了么？如果要釋放GPU的話(huà)，就需要考慮如何重新加載GPU。
所以，模型model和數(shù)據(jù)data不能放在項(xiàng)目啟動(dòng)的時(shí)候加載，只能放在調(diào)用訓(xùn)練的函數(shù)時(shí)加載，但是由于加載比較慢，所以只能放在一個(gè)異步的子線(xiàn)程或者子進(jìn)程中運(yùn)行。
所以，我這邊首先將模型數(shù)據(jù)的加載過(guò)程以及訓(xùn)練放在了一個(gè)多帶帶的線(xiàn)程中執(zhí)行。

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

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


輕點(diǎn)一下，即找到了答案。那就是在訓(xùn)練完成之后torch.cuda.empty_cache() 。代碼加上之后再運(yùn)行，發(fā)現(xiàn)并沒(méi)啥卵用?。。?！，CV大法第一運(yùn)用失敗
這到底是啥原因呢？我們后面會(huì)分析到?。?！

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

既然子線(xiàn)程加載模型并進(jìn)行訓(xùn)練不能釋放GPU的話(huà)，那么我們能不能轉(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 mpdef 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è)線(xiàn)程去執(zhí)行sub_process_train方法以保證訓(xùn)練接口可以正常返回。

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

代碼寫(xiě)好了，見(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。

這里有問(wèn)題，就是 forked 是啥，spawn 又是啥？這里就需要了解創(chuàng)建子進(jìn)程的方式了。
通過(guò)torch.multiprocessing.Process(target=training, args=(train_queue)) 創(chuàng)建一個(gè)子進(jìn)程
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。


一切看起來(lái)都很prefect。 But，But。通過(guò)gunicorn啟動(dòng)項(xiàng)目之后，再次調(diào)用接口，則出現(xiàn)下面結(jié)果。

用gunicorn啟動(dòng)項(xiàng)目子進(jìn)程竟然未執(zhí)行，這就很頭大了。不加mp.set_start_method(‘spawn’) 方法模型數(shù)據(jù)不能加載，
加上這個(gè)方法子進(jìn)程不能執(zhí)行，真的是一個(gè)頭兩個(gè)大。

第三階段（全局線(xiàn)程池+釋放GPU）

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

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

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

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

  pool.apply_async(func=async_produce_poets)

3. 用線(xiàn)程加載模型和釋放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()

粉絲專(zhuān)屬福利

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

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

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

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

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

											?? 驗(yàn)證碼 可通過(guò)搜索下方 公眾號(hào) 獲取?? 

?