Question

我非常了解要加载字典，然后要使用旧的参数字典（例如this great question & answer）加载实例。不幸的是，当我有一个torch.nn.Sequential时，我当然没有它的类定义。

所以我想仔细检查一下，什么是正确的方法。我相信torch.save就足够了（到目前为止我的代码还没有崩溃），尽管这些事情可能比人们预期的要微妙得多（例如，当我使用泡菜时会收到警告，但是torch.save在内部使用它，因此令人困惑）。另外，numpy拥有自己的保存功能（例如，参见this answer），这些保存功能往往效率更高，因此我可能忽略了一个微妙的取舍。

我的测试代码：


# creating data and running through a nn and saving it

import torch
import torch.nn as nn

from pathlib import Path
from collections import OrderedDict

import numpy as np

import pickle

path = Path('~/data/tmp/').expanduser()
path.mkdir(parents=True, exist_ok=True)

num_samples = 3
Din, Dout = 1, 1
lb, ub = -1, 1

x = torch.torch.distributions.Uniform(low=lb, high=ub).sample((num_samples, Din))

f = nn.Sequential(OrderedDict([
    ('f1', nn.Linear(Din,Dout)),
    ('out', nn.SELU())
]))
y = f(x)

# save data torch to numpy
x_np, y_np = x.detach().cpu().numpy(), y.detach().cpu().numpy()
np.savez(path / 'db', x=x_np, y=y_np)

print(x_np)
# save model
with open('db_saving_seq', 'wb') as file:
    pickle.dump({'f': f}, file)

# load model
with open('db_saving_seq', 'rb') as file:
    db = pickle.load(file)
    f2 = db['f']

# test that it outputs the right thing
y2 = f2(x)

y_eq_y2 = y == y2
print(y_eq_y2)

db2 = {'f': f, 'x': x, 'y': y}
torch.save(db2, path / 'db_f_x_y')

print('Done')

db3 = torch.load(path / 'db_f_x_y')
f3 = db3['f']
x3 = db3['x']
y3 = db3['y']
yy3 = f3(x3)

y_eq_y3 = y == y3
print(y_eq_y3)

y_eq_yy3 = y == yy3
print(y_eq_yy3)

如何将torch.nn.Sequential模型正确保存在pytorch中？

1 个答案: