10 · 把模型交出去：保存、TorchScript、ONNX、torch.export、compile

import torch
import torch.nn as nn


class Net(nn.Module):
    def __init__(self, hidden=8):
        super().__init__()
        self.fc1 = nn.Linear(4, hidden)
        self.fc2 = nn.Linear(hidden, 2)

    def forward(self, x):
        return self.fc2(torch.relu(self.fc1(x)))

import os
import torch
from mymodel import Net

torch.manual_seed(0)
m = Net()

torch.save(m, "whole.pt")                 # ① 整个模型对象
torch.save(m.state_dict(), "sd.pt")       # ② 只存权重

for f in ("whole.pt", "sd.pt"):
    print(f"{f:10} {os.path.getsize(f):6d} 字节")

print("state_dict 的 key:", list(m.state_dict().keys()))
print("类型:", type(m.state_dict()).__name__)

import io
import pickletools
import zipfile


def globals_in(path):
    with zipfile.ZipFile(path) as z:
        name = [n for n in z.namelist() if n.endswith("data.pkl")][0]
        blob = z.read(name)
    out = io.StringIO()
    pickletools.dis(blob, out)
    return sorted({ln.split("GLOBAL")[1].strip().strip("'")
                   for ln in out.getvalue().splitlines() if "GLOBAL" in ln})


print("whole.pt 里引用的类:")
for g in globals_in("whole.pt"):
    print("   ", g)
print("sd.pt 里引用的类:")
for g in globals_in("sd.pt"):
    print("   ", g)

import torch

print("--- ① 整个模型 ---")
try:
    m = torch.load("whole.pt", weights_only=False)
    print("加载成功:", type(m).__name__)
except Exception as e:
    print(type(e).__name__ + ": " + str(e))

print("--- ② 只有权重 ---")
sd = torch.load("sd.pt")
print("加载成功:", list(sd.keys()))

import torch
from mymodel import Net          # ⭐ 类定义在，也照样试试

m = torch.load("whole.pt")       # torch 2.13 的默认参数

import argparse
import numpy as np
import torch
from mymodel import Net

torch.manual_seed(0)
m = Net()
opt = torch.optim.Adam(m.parameters())
opt.step()

ckpt = {"model": m.state_dict(), "optim": opt.state_dict(), "epoch": 3, "best": 0.91}
torch.save(ckpt, "ckpt.pt")
back = torch.load("ckpt.pt")
print("① 标准 checkpoint:", list(back.keys()), " epoch =", back["epoch"])

torch.save({"cfg": argparse.Namespace(lr=1e-3), "sd": m.state_dict()}, "ckpt2.pt")
try:
    torch.load("ckpt2.pt")
except Exception as e:
    print("② 塞了个 Namespace ->", type(e).__name__)
    print("   ", [ln.strip() for ln in str(e).splitlines() if "Unsupported global" in ln])

torch.save({"arr": np.arange(3), "sd": m.state_dict()}, "ckpt3.pt")
try:
    torch.load("ckpt3.pt")
except Exception as e:
    print("③ 塞了个 numpy 数组 ->", type(e).__name__)
    print("   ", [ln.strip() for ln in str(e).splitlines() if "Unsupported global" in ln])

import torch
import torch.nn as nn


class Cond(nn.Module):
    def forward(self, x):
        if x.sum() > 0:
            return x * 2
        return x * -1


m = Cond()
pos = torch.tensor([1.0, 2.0])
neg = torch.tensor([-1.0, -2.0])

print("原模型      :", m(pos).tolist(), m(neg).tolist())

traced = torch.jit.trace(m, pos)          # ⚠️ 只喂了正数
print("trace 之后  :", traced(pos).tolist(), traced(neg).tolist())

scripted = torch.jit.script(m)
print("script 之后 :", scripted(pos).tolist(), scripted(neg).tolist())

print("--- trace 出来的代码 ---")
print(traced.code)
print("--- script 出来的代码 ---")
print(scripted.code)

import os
import torch
from mymodel import Net

torch.manual_seed(0)
m = Net().eval()
x = torch.randn(1, 4)
print("原模型输出:", [round(v, 4) for v in m(x)[0].tolist()])

torch.jit.save(torch.jit.script(m), "scripted.pt")
torch.save(x, "x.pt")
print("scripted.pt", os.path.getsize("scripted.pt"), "字节")

import torch                     # ⭐ 注意：这个目录里没有 mymodel.py

m = torch.jit.load("scripted.pt")
x = torch.load("x.pt")
print("类型:", type(m).__name__)
print("输出:", [round(v, 4) for v in m(x)[0].tolist()])
print("还能看到源码:")
print(m.code)

ModuleNotFoundError: No module named 'onnxscript'

DeprecationWarning: You are using the legacy TorchScript-based ONNX export. Starting in PyTorch 2.9, the new torch.export-based ONNX exporter has become the default.

import os
import torch
from mymodel import Net

torch.manual_seed(0)
m = Net().eval()
x = torch.randn(1, 4)

torch.onnx.export(
    m, (x,), "net.onnx",
    input_names=["x"], output_names=["y"],
    dynamic_axes={"x": {0: "batch"}, "y": {0: "batch"}},   # ⭐ 声明哪一维是可变的
    dynamo=False,
)
print("net.onnx", os.path.getsize("net.onnx"), "字节")
print("原模型输出:", [round(v, 4) for v in m(x)[0].tolist()])
torch.save(x, "x1.pt")

import numpy as np
import onnx
import onnxruntime as ort
import torch

g = onnx.load("net.onnx").graph
print("输入 :", [(i.name, [d.dim_param or d.dim_value
                          for d in i.type.tensor_type.shape.dim]) for i in g.input])
print("输出 :", [(o.name, [d.dim_param or d.dim_value
                          for d in o.type.tensor_type.shape.dim]) for o in g.output])
print("节点 :", [n.op_type for n in g.node])
print("权重 :", [(w.name, list(w.dims)) for w in g.initializer])

sess = ort.InferenceSession("net.onnx")
x = torch.load("x1.pt")
print("onnxruntime 输出:", [round(v, 4) for v in sess.run(None, {"x": x.numpy()})[0][0].tolist()])

big = np.random.randn(5, 4).astype(np.float32)
print("换成 batch=5 :", sess.run(None, {"x": big})[0].shape)

import onnx
import onnxruntime as ort
import torch
import torch.nn as nn


class Cond(nn.Module):
    def forward(self, x):
        if x.sum() > 0:
            return x * 2
        return x * -1


m = Cond().eval()
pos = torch.tensor([[1.0, 2.0]])
torch.onnx.export(m, (pos,), "cond.onnx", input_names=["x"], output_names=["y"],
                  dynamic_axes={"x": {0: "b"}, "y": {0: "b"}}, dynamo=False)
print("节点:", [n.op_type for n in onnx.load("cond.onnx").graph.node])

sess = ort.InferenceSession("cond.onnx")
print("正数 pytorch:", m(pos).tolist(), " onnx:", sess.run(None, {"x": pos.numpy()})[0].tolist())
neg = torch.tensor([[-1.0, -2.0]])
print("负数 pytorch:", m(neg).tolist(), " onnx:", sess.run(None, {"x": neg.numpy()})[0].tolist())

import numpy as np
import onnx
import onnxruntime as ort
import torch
import torch.nn as nn


class NpSqrt(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x):
        out = torch.from_numpy(np.sqrt(x.detach().numpy()))
        ctx.save_for_backward(out)
        return out

    @staticmethod
    def backward(ctx, g):
        (out,) = ctx.saved_tensors
        return g / (2 * out)


class M(nn.Module):
    def forward(self, x):
        return NpSqrt.apply(x)


m = M().eval()
x = torch.tensor([[1.0, 4.0, 9.0]])
torch.onnx.export(m, (x,), "np.onnx", dynamo=False)

g = onnx.load("np.onnx").graph
print("图里的节点:", [n.op_type for n in g.node])
print("图的输入  :", [i.name for i in g.input])
sess = ort.InferenceSession("np.onnx")
print("ort 要的输入:", [i.name for i in sess.get_inputs()])
print("跑一次    :", sess.run(None, {})[0].tolist())
print("pytorch 换个输入:", m(torch.tensor([[16.0, 25.0, 36.0]])).tolist())

import os
import torch
from mymodel import Net

torch.manual_seed(0)
m = Net().eval()
x = torch.randn(1, 4)

ep = torch.export.export(m, (x,))
print("类型:", type(ep).__name__)
print(ep.graph_module.code)
torch.export.save(ep, "net.pt2")
print("net.pt2", os.path.getsize("net.pt2"), "字节")
print("原模型:", [round(v, 4) for v in m(x)[0].tolist()])

import torch                      # ⭐ 这个目录里没有 mymodel.py

ep = torch.export.load("net.pt2")
x = torch.load("x.pt")        # ④ 存的那份输入，和 net.pt2 一起拷过来
print("类型:", type(ep).__name__)
print("输出:", [round(v, 4) for v in ep.module()(x)[0].tolist()])

import torch
import torch.nn as nn


class Cond(nn.Module):
    def forward(self, x):
        if x.sum() > 0:
            return x * 2
        return x * -1


torch.export.export(Cond().eval(), (torch.tensor([[1.0, 2.0]]),))

GuardOnDataDependentSymNode: Could not guard on data-dependent expression Eq(u0, 1) (unhinted: Eq(u0, 1)).  (Size-like symbols: none)

consider using data-dependent friendly APIs such as guard_or_false, guard_or_true and statically_known_true.

import torch
from mymodel import Net

torch.manual_seed(0)
m = Net().eval()
x2 = torch.randn(2, 4)
x5 = torch.randn(5, 4)

ep = torch.export.export(m, (x2,))                       # 只给了 batch=2 的样例
try:
    print("① 不声明动态维，喂 batch=5:", ep.module()(x5).shape)
except Exception as e:
    print("① 不声明动态维，喂 batch=5 ->", type(e).__name__ + ": " + str(e).strip()[:220])

b = torch.export.Dim("batch")
ep2 = torch.export.export(m, (x2,), dynamic_shapes={"x": {0: b}})
print("② 声明 batch 是动态维，喂 batch=5:", ep2.module()(x5).shape)

import torch
from mymodel import Net

torch.manual_seed(0)
m = Net().eval()
c = torch.compile(m)
print("torch.compile 返回的类型:", type(c).__name__)
print("它还是不是 nn.Module :", isinstance(c, torch.nn.Module))
print("state_dict 的 key    :", list(c.state_dict().keys()))
print("拿得到原模型吗       :", type(c._orig_mod).__name__)

import torch
from mymodel import Net

torch.manual_seed(0)
m = Net()
c = torch.compile(m)

torch.save(c.state_dict(), "compiled_sd.pt")     # ⚠️ 存的是编译后那个壳的 state_dict

fresh = Net()
fresh.load_state_dict(torch.load("compiled_sd.pt"))

RuntimeError:
Error(s) in loading state_dict for Net:
    Missing key(s) in state_dict: "fc1.weight", "fc1.bias", "fc2.weight", "fc2.bias". 
    Unexpected key(s) in state_dict: "_orig_mod.fc1.weight", "_orig_mod.fc1.bias", "_orig_mod.fc2.weight", "_orig_mod.fc2.bias". 

# 🧩 骨架：接着上一段跑，Net / c / torch 都在那里定义
torch.save(c._orig_mod.state_dict(), "ok_sd.pt")     # ⭐ 穿过那层壳
fresh2 = Net()
print(fresh2.load_state_dict(torch.load("ok_sd.pt")))

InductorError: InvalidCxxCompiler: Compiler: cl is not found.