Python 實作輕量級任務隊列優化繪圖效能

在使用梯度下降算法訓練神經網路並即時繪圖時，同步呼叫 plotly.write 會阻塞主線程，嚴重影響訓練速度。本文分享一個純 Python 實作的輕量級任務隊列，將繪圖任務卸載到背景線程，無需引入 Celery、Redis 等重型框架。

問題背景

梯度下降迭代過程中需頻繁繪圖，若每次呼叫 plotly.write 都同步等待網路回應，會顯著拖慢演算法執行效率。開啟新線程雖可緩解，但管理多線程較繁瑣；引入分散式任務隊列又過於笨重。

解決方案：自製 TaskQueue

繼承 Queue.Queue，建立固定數量的 Worker 線程（Daemon 模式），主線程只需 add_task 丟入任務，即可立即繼續執行。

核心代碼結構

from threading import Thread
import Queue
import time

class TaskQueue(Queue.Queue):
    def __init__(self, num_workers=1):
        Queue.Queue.__init__(self)
        self.num_workers = num_workers
        self.start_workers()

    def add_task(self, task, *args, **kwargs):
        args = args or ()
        kwargs = kwargs or {}
        self.put((task, args, kwargs))

    def start_workers(self):
        for i in range(self.num_workers):
            t = Thread(target=self.worker)
            t.daemon = True
            t.start()

    def worker(self):
        while True:
            task, args, kwargs = self.get()
            task(*args, **kwargs)
            self.task_done()

使用範例

def blokkah(*args, **kwargs):
    time.sleep(5)
    print("Blokkah mofo!")

q = TaskQueue(num_workers=5)
for item in range(10):
    q.add_task(blokkah)
q.join()  # 等待所有任務完成
print("All done!")

整合到梯度下降

def gradient_descent():
    # the gradient descent code
    queue.add_task(plotly.write, x=X, y=Y)

修改後，繪圖任務非同步執行，主線程無阻塞，訓練效率明顯提升。完整代碼已整合 TaskQueue 類與測試函數，可直接複製使用。