您好，我是碼農(nóng)飛哥，感謝您閱讀本文，歡迎一鍵三連哦。
本文只是記錄我優(yōu)化的心酸歷程。無他，唯記錄爾。。。。。小伙伴們可圍觀，可打call，可以私信與我交流。
干貨滿滿，建議收藏，需要用到時?？纯?。 小伙伴們?nèi)缬袉栴}及需要，歡迎踴躍留言哦~ ~ ~。

問題背景

現(xiàn)有一個古詩自動生成的訓(xùn)練接口，該接口通過Pytorch來生訓(xùn)練模型（即生成古詩）為了加速使用到了GPU，但是訓(xùn)練完成之后GPU未能釋放。故此需要進(jìn)行優(yōu)化，即在古詩生成完成之后釋放GPU。
該項目是一個通過Flask搭建的web服務(wù)，在服務(wù)器上為了實現(xiàn)并發(fā)采用的是gunicorn來啟動應(yīng)用。通過pythorch來進(jìn)行古詩訓(xùn)練。項目部署在一個CentOS的服務(wù)器上。

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

因為特殊的原因這里之后一個服務(wù)器供使用，故不能考慮加機(jī)器的情況。

優(yōu)化歷程

pytorch在訓(xùn)練模型時，需要先加載模型model和數(shù)據(jù)data，如果有GPU顯存的話我們可以將其放到GPU顯存中加速，如果沒有GPU的話則只能使用CPU了。
由于加載模型以及數(shù)據(jù)的過程比較慢。所以，我這邊將加載過程放在了項目啟動時加載。

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()

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

小小分析一波

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

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

GPU沒釋放，那就釋放唄。這不是很簡單么？百度一波pytorch怎么釋放GPU顯存。


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

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

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

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

def training(queue):    manage.app.app_context().push()    current_app.logger.error("基礎(chǔ)加載開始")    with manage.app.app_context():        device = "cuda" if torch.cuda.is_available() else "cpu"        current_app.logger.error("device1111開始啦啦啦")        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é)果放入隊列中        queue.put(result_list)

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

from torch import multiprocessing as mpdef sub_process_train():	#定義一個隊列獲取訓(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. 因為子進(jìn)程要阻塞獲取執(zhí)行結(jié)果，所以需要定義一個線程去執(zhí)行sub_process_train方法以保證訓(xùn)練接口可以正常返回。

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

代碼寫好了，見證奇跡的時候來了。
首先用python manage.py 啟動一下，看下結(jié)果，運(yùn)行結(jié)果如下，報了一個錯誤，從錯誤的提示來看就是不能在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。

這里有問題，就是 forked 是啥，spawn 又是啥？這里就需要了解創(chuàng)建子進(jìn)程的方式了。
通過torch.multiprocessing.Process(target=training, args=(train_queue)) 創(chuàng)建一個子進(jìn)程
fork和spawn是構(gòu)建子進(jìn)程的不同方式，區(qū)別在于
1. fork: 除了必要的啟動資源，其余的變量，包，數(shù)據(jù)等都集成自父進(jìn)程，也就是共享了父進(jìn)程的一些內(nèi)存頁，因此啟動較快，但是由于大部分都是用的自父進(jìn)程數(shù)據(jù)，所有是不安全的子進(jìn)程。
2. spawn：從頭構(gòu)建一個子進(jìn)程，父進(jìn)程的數(shù)據(jù)拷貝到子進(jìn)程的空間中，擁有自己的Python解釋器，所有需要重新加載一遍父進(jìn)程的包，因此啟動叫慢，但是由于數(shù)據(jù)都是自己的，安全性比較高。
回到剛剛那個報錯上面去。為啥提示要不能重復(fù)加載。
這是因為Python3中使用 spawn啟動方法才支持在進(jìn)程之間共享CUDA張量。而用的multiprocessing 是使用 fork 創(chuàng)建子進(jìn)程，不被 CUDA 運(yùn)行時所支持。
所以，只有在創(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))    ##省略其他代碼

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


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

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

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

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

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

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

2. 通過線程池來執(zhí)行訓(xùn)練

  pool.apply_async(func=async_produce_poets)

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

def async_produce_poets():    try:        print("子進(jìn)程開始" + 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)練總耗時是={0}".format(str(train_seconds)))    except Exception as e:        manage.app.app_context().push()

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

這里因為使用到了gunicorn來啟動項目。所以gunicorn 相關(guān)的知識必不可少。在CPU受限的系統(tǒng)中采用sync的工作模式比較理想。
詳情可以查看gunicorn的簡單總結(jié)

問題分析，前面第一階段直接使用torch.cuda.empty_cache() 沒能釋放GPU就是因為沒有刪除掉模型model。模型已經(jīng)加載到了GPU了。

總結(jié)

本文從實際項目的優(yōu)化入手，記錄優(yōu)化方面的方方面面。希望對讀者朋友們有所幫助。

參考

multiprocessing fork() vs spawn()

粉絲專屬福利

軟考資料：實用軟考資料

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

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

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

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

											?? 驗證碼 可通過搜索下方 公眾號 獲取?? 

?