記一次效能優化的心酸歷程【Flask+Gunicorn+pytorch+多程式+執行緒池，一頓操作猛如虎】

碼農飛哥發表於2021-11-13
優化FlaskPyTorch執行緒

您好，我是碼農飛哥，感謝您閱讀本文，歡迎一鍵三連哦。
本文只是記錄我優化的心酸歷程。無他，唯記錄爾。。。。。小夥伴們可圍觀，可打call，可以私信與我交流。
乾貨滿滿，建議收藏，需要用到時常看看。 小夥伴們如有問題及需要，歡迎踴躍留言哦~ ~ ~。

問題背景

現有一個古詩自動生成的訓練介面，該介面通過Pytorch來生訓練模型（即生成古詩）為了加速使用到了GPU，但是訓練完成之後GPU未能釋放。故此需要進行優化，即在古詩生成完成之後釋放GPU。
該專案是一個通過Flask搭建的web服務，在伺服器上為了實現併發採用的是gunicorn來啟動應用。通過pythorch來進行古詩訓練。專案部署在一個CentOS的伺服器上。

系統環境

軟體 版本
flask 0.12.2
gunicorn 19.9.0
CentOS 6.6 帶有GPU的伺服器，不能加機器
pytorch 1.7.0+cpu

因為特殊的原因這裡之後一個伺服器供使用，故不能考慮加機器的情況。

優化歷程

pytorch在訓練模型時，需要先載入模型model和資料data，如果有GPU視訊記憶體的話我們可以將其放到GPU視訊記憶體中加速，如果沒有GPU的話則只能使用CPU了。
由於載入模型以及資料的過程比較慢。所以，我這邊將載入過程放在了專案啟動時載入。

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

這部分耗時大約在6秒左右。cuda表示使用torch的cuda。模型資料載入之後所佔的GPU視訊記憶體大小大約在1370MB。優化的目標就是在訓練完成之後將這部分佔用的視訊記憶體釋放掉。

小小分析一波

現狀是專案啟動時就載入模型model和資料data的話，當模型資料在GPU中釋放掉之後，下次再進行模型訓練的話不就沒有模型model和資料data了麼？如果要釋放GPU的話，就需要考慮如何重新載入GPU。
所以，模型model和資料data不能放在專案啟動的時候載入，只能放在呼叫訓練的函式時載入，但是由於載入比較慢，所以只能放在一個非同步的子執行緒或者子程式中執行。
所以，我這邊首先將模型資料的載入過程以及訓練放在了一個單獨的執行緒中執行。

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

GPU沒釋放，那就釋放唄。這不是很簡單麼？百度一波pytorch怎麼釋放GPU視訊記憶體。
在這裡插入圖片描述
在這裡插入圖片描述
輕點一下，即找到了答案。那就是在訓練完成之後torch.cuda.empty_cache() 。程式碼加上之後再執行，發現並沒啥卵用！！！！，CV大法第一運用失敗
這到底是啥原因呢？我們後面會分析到！！！

第二階段（建立子程式載入模型並進行訓練）

既然子執行緒載入模型並進行訓練不能釋放GPU的話，那麼我們能不能轉變一下思路。建立一個子程式來載入模型資料並進行訓練，
當訓練完成之後就將這個子程式殺掉，它所佔用的資源（主要是GPU視訊記憶體）不就被釋放了麼？
這思路看起來沒有絲毫的毛病呀。說幹就幹。

  1. 定義載入模型資料以及訓練的方法 training。（程式碼僅供參考）
def training(queue):
    manage.app.app_context().push()
    current_app.logger.error('基礎載入開始')
    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('基礎載入完成')
		#訓練方法
		result_list=start_train(model,n_ctx,device)
        current_app.logger.error('完成訓練')
        #將訓練方法返回的結果放入佇列中
        queue.put(result_list)
  1. 建立子程式執行training方法，然後通過阻塞的方法獲取訓練結果
from torch import multiprocessing as mp
def sub_process_train():
	#定義一個佇列獲取訓練結果
    train_queue = mp.Queue()
    training_process = mp.Process(target=training, args=(train_queue))
    training_process.start()
    current_app.logger.error('子程式執行')
    # 等訓練完成
    training_process.join()
    current_app.logger.error('執行完成')
    #獲取訓練結果
    result_list = train_queue.get()
    current_app.logger.error('獲取到資料')
    if training_process.is_alive():
        current_app.logger.error('子程式還存活')
        #殺掉子程式
        os.kill(training_process.pid, signal.SIGKILL)
    current_app.logger.error('殺掉子程式')
    return result_list
  1. 因為子程式要阻塞獲取執行結果，所以需要定義一個執行緒去執行sub_process_train方法以保證訓練介面可以正常返回。
import threading
threading.Thread(target=sub_process_train).start()

程式碼寫好了，見證奇蹟的時候來了。

首先用python manage.py 啟動一下，看下結果，執行結果如下，報了一個錯誤，從錯誤的提示來看就是不能在forked的子程式中重複載入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 又是啥？這裡就需要了解建立子程式的方式了。
通過torch.multiprocessing.Process(target=training, args=(train_queue)) 建立一個子程式

fork和spawn是構建子程式的不同方式，區別在於
1. fork: 除了必要的啟動資源，其餘的變數，包，資料等都整合自父程式，也就是共享了父程式的一些記憶體頁，因此啟動較快，但是由於大部分都是用的自父程式資料，所有是不安全的子程式。
2. spawn：從頭構建一個子程式，父程式的資料拷貝到子程式的空間中，擁有自己的Python直譯器，所有需要重新載入一遍父程式的包，因此啟動叫慢，但是由於資料都是自己的，安全性比較高。

回到剛剛那個報錯上面去。為啥提示要不能重複載入。

這是因為Python3中使用 spawn啟動方法才支援在程式之間共享CUDA張量。而用的multiprocessing 是使用 fork 建立子程式，不被 CUDA 執行時所支援。
所以，只有在建立子程式之前加上 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 執行專案。執行結果圖1和圖2所示，可以看出可以正確是使用GPU視訊記憶體，在訓練完成之後也可以釋放GPU。
圖1
圖2
一切看起來都很prefect。 But，But。通過gunicorn啟動專案之後，再次呼叫介面，則出現下面結果。
在這裡插入圖片描述
用gunicorn啟動專案子程式竟然未執行，這就很頭大了。不加mp.set_start_method('spawn') 方法模型資料不能載入，
加上這個方法子程式不能執行，真的是一個頭兩個大。

第三階段（全域性執行緒池+釋放GPU）

子程式的方式也不行了。只能回到前面的執行緒方式了。前面建立執行緒的方式都是直接通過直接new一個新執行緒的方式，當同時執行的執行緒數過多的話，則很容易就會出現GPU佔滿的情況，從而導致應用崩潰。所以，這裡採用全域性執行緒池的方式來建立並管理執行緒，然後當執行緒執行完成之後釋放資源。

  1. 在專案啟動之後就建立一個全域性執行緒池。大小是2。保證還有剩餘的GPU。
from multiprocessing.pool import ThreadPool
pool = ThreadPool(processes=2)
  1. 通過執行緒池來執行訓練
  pool.apply_async(func=async_produce_poets)
  1. 用執行緒載入模型和釋放GPU

def async_produce_poets():
    try:
        print("子程式開始" + 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轉到cpu
        model = model.to('cpu')
        #刪除模型，也就是刪除引用
        del model
        #在使用其釋放GPU。
        torch.cuda.empty_cache()
        train_seconds = int(time.time() - start_time)
        current_app.logger.info('訓練總耗時是={0}'.format(str(train_seconds)))
    except Exception as e:
        manage.app.app_context().push()

這一番操作之後，終於達到了理想的效果。

這裡因為使用到了gunicorn來啟動專案。所以gunicorn 相關的知識必不可少。在CPU受限的系統中採用sync的工作模式比較理想。
詳情可以檢視gunicorn的簡單總結

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

總結

本文從實際專案的優化入手，記錄優化方面的方方面面。希望對讀者朋友們有所幫助。

參考

multiprocessing fork() vs spawn()

粉絲專屬福利

軟考資料：實用軟考資料

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

學習資料：50G的各類學習資料

脫單祕籍：回覆【脫單】

併發程式設計：回覆【併發程式設計】

全網同名【碼農飛哥】。不積跬步，無以至千里，享受分享的快樂
我是碼農飛哥，再次感謝您讀完本文。

相關文章