如何通过Flask端点高效返回PyTorch模型state_dict权重？

核心问题分析 #

直接将state_dict转为列表再以JSON返回的方式，会触发大量张量到Python列表的逐元素转换，对大模型来说速度极慢，且JSON格式体积大、传输效率低。最优方案是利用PyTorch原生的二进制序列化机制，结合Flask的send_file直接返回二进制权重文件。

解决方案1：基于torch.save的二进制流返回 #

利用torch.save将state_dict序列化到内存字节流，再通过send_file返回，无需磁盘IO，效率极高。接收端可直接用torch.load加载，无需额外转换。

from flask import Flask, send_file
import torch
from io import BytesIO

app = Flask(__name__)
# 假设model是已加载的PyTorch模型
model = torch.nn.Linear(10, 2)

@app.route('/send_weights', methods=['GET', 'POST'])
def send_weights():
    # 创建内存字节流
    buffer = BytesIO()
    # 将state_dict序列化到字节流
    torch.save(model.state_dict(), buffer)
    # 将指针移到流的开头
    buffer.seek(0)
    # 返回二进制文件，设置正确的mimetype和下载文件名
    return send_file(
        buffer,
        mimetype='application/octet-stream',
        as_attachment=True,
        download_name='model_weights.pt'
    )

接收端处理示例（Python） #

如果接收方也是Python环境，直接用torch.load加载响应内容即可：

import requests
import torch
from io import BytesIO

response = requests.get('http://your-server:port/send_weights')
weights = torch.load(BytesIO(response.content))
# 加载到模型
model.load_state_dict(weights)

解决方案2：可选压缩优化（针对大模型） #

如果模型权重很大，可以在序列化后添加gzip压缩，进一步减少传输体积：

from flask import Flask, send_file
import torch
from io import BytesIO
import gzip

app = Flask(__name__)
model = torch.nn.Linear(10, 2)

@app.route('/send_weights', methods=['GET', 'POST'])
def send_weights():
    buffer = BytesIO()
    # 先序列化到字节流，再压缩
    with gzip.GzipFile(fileobj=buffer, mode='wb') as gz:
        torch.save(model.state_dict(), gz)
    buffer.seek(0)
    return send_file(
        buffer,
        mimetype='application/gzip',
        as_attachment=True,
        download_name='model_weights.pt.gz'
    )

接收端解压后加载：

import requests
import torch
from io import BytesIO
import gzip

response = requests.get('http://your-server:port/send_weights')
with gzip.GzipFile(fileobj=BytesIO(response.content), mode='rb') as gz:
    weights = torch.load(gz)
model.load_state_dict(weights)

为什么这比转列表高效？ #

1. 序列化效率：torch.save是PyTorch原生的二进制序列化，直接操作张量的底层数据，避免了Python列表的逐元素转换开销。
2. 传输效率：二进制文件体积远小于JSON格式的列表，减少带宽占用和传输时间。
3. 接收端易用性：无需手动将列表转回张量，直接用torch.load即可恢复state_dict。

内容的提问来源于stack exchange，提问作者amlarraz