Một gói suy luận cần cả chuẩn hóa và bảng nhãn, không chỉ state_dict

import io
import torch

model = torch.nn.Linear(2, 2, bias=False)
with torch.no_grad():
    model.weight.copy_(torch.tensor([[1., 0.], [-1., 0.]]))
artifact = {'version': 1, 'state_dict': model.state_dict(), 'mean': torch.tensor([10., 0.]), 'std': torch.tensor([2., 1.]), 'labels': ['A', 'B']}
stream = io.BytesIO()
torch.save(artifact, stream)
stream.seek(0)
loaded = torch.load(stream, map_location='cpu', weights_only=True)
restored = torch.nn.Linear(2, 2, bias=False)
restored.load_state_dict(loaded['state_dict'])
restored.eval()
raw = torch.tensor([[9., 0.], [11., 0.]])
assert torch.all(loaded['std'] > 0)
with torch.inference_mode():
    correct = restored((raw-loaded['mean'])/loaded['std']).argmax(1)
    wrong = restored(raw).argmax(1)
expected = torch.tensor([1, 0])
assert torch.equal(correct, expected)
assert loaded['labels'] == ['A', 'B'] and loaded['version'] == 1
result = {'Dung hop dong tien xu ly': int((correct != expected).sum()), 'Bo qua chuan hoa': int((wrong != expected).sum())}
chart_labels, chart_values = list(result), list(result.values())
chart_unit = 'So du doan khac ket qua tham chieu'
print(result, 'labels:', [loaded['labels'][i] for i in correct.tolist()])