Pytorch自定义层含lambda函数（Lambda层）时保存模型出错的解决方案

出错

类型：

_pickle.PicklingError: Can't pickle lambda> at 0x0000000018944488>: attribute lookup <lambda> on __main__ failed

详情：

File "C:\Users\piantou\AppData\Local\Programs\Python\Python36\lib\site-packages\torch\serialization.py", line 224, in save
    return _with_file_like(f, "wb", lambda f: _save(obj, f, pickle_module, pickle_protocol))
  File "C:\Users\piantou\AppData\Local\Programs\Python\Python36\lib\site-packages\torch\serialization.py", line 149, in _with_file_like
    return body(f)
  File "C:\Users\piantou\AppData\Local\Programs\Python\Python36\lib\site-packages\torch\serialization.py", line 224, in 
    return _with_file_like(f, "wb", lambda f: _save(obj, f, pickle_module, pickle_protocol))
  File "C:\Users\piantou\AppData\Local\Programs\Python\Python36\lib\site-packages\torch\serialization.py", line 297, in _save
    pickler.dump(obj)
_pickle.PicklingError: Can't pickle  at 0x0000000018944488>: attribute lookup  on __main__ failed

原因

以Mnist数据集识别为例，在建立的网络模型中自定义了Lambda层（使用lambda函数），而save不能pickle lambda函数。具体代码如下：

import torch
import torch.nn as nn
import torch.utils.data as Data
import torchvision  # 数据库
import torch.nn.functional as F

torch.manual_seed(1)

epoches = 1  # 训练整批数据的次数
batch_size = 500  # 每次训练的样本数
lr = 0.001

# 1.获得训练数据集
train_data = torchvision.datasets.MNIST(
    root='./mnist',  # 下载路径
    train=True,  # 下载训练数据
    transform=torchvision.transforms.ToTensor(),  # 将数据转化为tensor类型
    download=False  # 是否下载MNIST数据集
)
# 获得测试数据集
test_data = torchvision.datasets.MNIST(
    root='./mnist',
    train=False
)

# 将dataset放入DataLoader中  (batch, channel, 28, 28)
loader = Data.DataLoader(
    dataset=train_data,
    batch_size=batch_size,  # 设置batch size
    shuffle=True  # 打乱数据
)

# 利用前20个样本和标签测试，并归一化
test_x = torch.unsqueeze(test_data.test_data, dim=1).type(
    torch.FloatTensor)[:20] / 255
test_y = test_data.test_labels[:20]


# 2.前向传播过程

class Lambda(nn.Module):
    def __init__(self, func):
        super().__init__()
        self.func = func

    def forward(self, x):
        return self.func(x)


model=nn.Sequential(
    nn.Conv2d(
                in_channels=1,  # 输入图片深度
                out_channels=16,  # 滤波器数量
                kernel_size=5,  # 滤波器大小
                stride=1,  # 步长
                padding=2),  # 保持输出图片大小不变 padding=(kernel_size-stride)/2 )
    nn.ReLU(),
    nn.MaxPool2d(kernel_size=2),
    nn.Conv2d(16, 32, 3, 1, 1),
    nn.ReLU(),
    nn.MaxPool2d(2),
    Lambda(lambda x: x.view(x.size(0), -1)),  #此处使用了自定义Lambda层
    nn.Linear(32 * 7 * 7, 10)
)


# 训练模型的同时保存网络模型参数
def save():
    # 3.利用自定义的前向传播过程设计网络
    cnn = model

    # 4.设置优化算法、学习率
    optimizer = torch.optim.Adam(cnn.parameters(), lr=lr)

    # 5.设置损失函数
    loss_func = torch.nn.CrossEntropyLoss()

    # 6.迭代训练
    for epoch in range(epoches):
        for step, (batch_x, batch_y) in enumerate(loader):
            out = cnn(batch_x)  # 输入训练集，获得当前迭代输出值
            loss = loss_func(out, batch_y)  # 获得当前迭代的损失

            optimizer.zero_grad()  # 清除上次迭代的更新梯度
            loss.backward()  # 反向传播
            optimizer.step()  # 更新权重

            # 打印训练过程中的测试集准确率
            if step % 500 == 0:
                test_out = cnn(test_x)  # 输入测试集
                # 获得当前softmax层最大概率对应的索引值
                pred = torch.max(test_out, 1)[1]
                # 将二维压缩为一维
                pred_y = pred.data.numpy().squeeze()
                label_y = test_y.data.numpy()
                accuracy = sum(pred_y == label_y) / test_y.size()
                print("第 %d 个epoch，第 %d 次迭代，准确率为 %.2f" % (epoch + 1, step / 100 + 1, accuracy))

    # 7.保存模型结构和参数
    torch.save(cnn, 'cnn.pth')
    torch.save(cnn.state_dict(), 'cnn_statedict.pth')
    for i in model.parameters():
        print(i)
if __name__ == '__main__':
    save()

保存模型结构和参数这一句出错：

 torch.save(cnn, 'cnn.pth')

而仅仅保存模型参数是不出错的：

torch.save(cnn.state_dict(), 'cnn_statedict.pth')

所以很多所谓的解决方案中就说了：你用torch.save(modeltosave.state_dict(), 'model_statedict.pth')就好了。too young too simple!!

 

解决方案

参考《[Solved] Using a fastai-trained model with plain Pytorch》

（1）pip install dill后：

import dill

（2） 将原代码： torch.save(cnn, 'cnn.pkl')改为：

torch.save(cnn, 'cnn.pth',pickle_module=dill)

后续的网络模型加载

pytorch与保存、加载模型有关的常用函数3个：

• torch.save(): 保存一个序列化的对象到磁盘，使用的是Python的pickle库来实现的。
• torch.load(): 反序列化一个pickled对象并加载到内存当中。
• torch.nn.Module.load_state_dict(): 加载一个反序列化的state_dict对象

 

（1）加载整个网络（加载前必须定义好模型结构）

重要一句是（注意加pickle_module=dill）：

model=torch.load('cnn.pkl',pickle_module=dill)

全部代码： 

import torch
import dill
import torch.nn as nn
import torch.utils.data as Data
import torchvision  # 数据库
import torch.nn.functional as F

torch.manual_seed(1)

epoches = 1  # 训练整批数据的次数
batch_size = 500  # 每次训练的样本数
lr = 0.001

# 1.获得训练数据集
train_data = torchvision.datasets.MNIST(
    root='./mnist',  # 下载路径
    train=True,  # 下载训练数据
    transform=torchvision.transforms.ToTensor(),  # 将数据转化为tensor类型
    download=False  # 是否下载MNIST数据集
)
# 获得测试数据集
test_data = torchvision.datasets.MNIST(
    root='./mnist',
    train=False
)

# 将dataset放入DataLoader中  (batch, channel, 28, 28)
loader = Data.DataLoader(
    dataset=train_data,
    batch_size=batch_size,  # 设置batch size
    shuffle=True  # 打乱数据
)

# 利用前20个样本和标签测试，并归一化
test_x = torch.unsqueeze(test_data.test_data, dim=1).type(
    torch.FloatTensor)[:20] / 255
test_y = test_data.test_labels[:20]

class Lambda(nn.Module):
    def __init__(self, func):
        super().__init__()
        self.func = func

    def forward(self, x):
        return self.func(x)


model=nn.Sequential(
    nn.Conv2d(
                in_channels=1,  # 输入图片深度
                out_channels=16,  # 滤波器数量
                kernel_size=5,  # 滤波器大小
                stride=1,  # 步长
                padding=2),  # 保持输出图片大小不变 padding=(kernel_size-stride)/2 )
    nn.ReLU(),
    nn.MaxPool2d(kernel_size=2),
    nn.Conv2d(16, 32, 3, 1, 1),
    nn.ReLU(),
    nn.MaxPool2d(2),
    Lambda(lambda x: x.view(x.size(0), -1)),
    nn.Linear(32 * 7 * 7, 10)
)


model=torch.load('cnn.pkl',pickle_module=dill)


def eval():
    model.eval()
    with torch.no_grad():
        test_out = model(test_x)  # 输入测试集
        # 获得当前softmax层最大概率对应的索引值
        pred = torch.max(test_out, 1)[1]
        # 将二维压缩为一维
        pred_y = pred.data.numpy().squeeze()
        label_y = test_y.data.numpy()
        accuracy = sum(pred_y == label_y) / test_y.size()
        print("准确率为 %.2f" % (accuracy))

if __name__ == '__main__':
    eval()

（2）加载cnn_state

重要代码是：

model.load_state_dict(torch.load('cnn.pth'))

注意：

请注意 load_state_dict() 函数只接受字典对象，而不是保存对象的路径。这就意味着在你传给 load_state_dict() 函数之前，你必须反序列化（通过torch.load()实现）你保存的 state_dict。例如，你无法通过 model.load_state_dict(PATH)来加载模型。在 Pytorch 中最常见的模型保存使用 ‘.pt’ 或者是 ‘.pth’ 作为模型文件PATH的扩展名。

 

全部代码：

import torch
import dill
import torch.nn as nn
import torch.utils.data as Data
import torchvision  # 数据库
import torch.nn.functional as F

torch.manual_seed(1)

epoches = 1  # 训练整批数据的次数
batch_size = 500  # 每次训练的样本数
lr = 0.001

# 1.获得训练数据集
train_data = torchvision.datasets.MNIST(
    root='./mnist',  # 下载路径
    train=True,  # 下载训练数据
    transform=torchvision.transforms.ToTensor(),  # 将数据转化为tensor类型
    download=False  # 是否下载MNIST数据集
)
# 获得测试数据集
test_data = torchvision.datasets.MNIST(
    root='./mnist',
    train=False
)

# 将dataset放入DataLoader中  (batch, channel, 28, 28)
loader = Data.DataLoader(
    dataset=train_data,
    batch_size=batch_size,  # 设置batch size
    shuffle=True  # 打乱数据
)

# 利用前20个样本和标签测试，并归一化
test_x = torch.unsqueeze(test_data.test_data, dim=1).type(
    torch.FloatTensor)[:20] / 255
test_y = test_data.test_labels[:20]

class Lambda(nn.Module):
    def __init__(self, func):
        super().__init__()
        self.func = func

    def forward(self, x):
        return self.func(x)


model=nn.Sequential(
    nn.Conv2d(
                in_channels=1,  # 输入图片深度
                out_channels=16,  # 滤波器数量
                kernel_size=5,  # 滤波器大小
                stride=1,  # 步长
                padding=2),  # 保持输出图片大小不变 padding=(kernel_size-stride)/2 )
    nn.ReLU(),
    nn.MaxPool2d(kernel_size=2),
    nn.Conv2d(16, 32, 3, 1, 1),
    nn.ReLU(),
    nn.MaxPool2d(2),
    Lambda(lambda x: x.view(x.size(0), -1)),
    nn.Linear(32 * 7 * 7, 10)
)


model.load_state_dict(torch.load('cnn.pth'))


def eval():
    # model.eval()
    with torch.no_grad():
        test_out = model(test_x)  # 输入测试集
        # 获得当前softmax层最大概率对应的索引值
        pred = torch.max(test_out, 1)[1]
        # 将二维压缩为一维
        pred_y = pred.data.numpy().squeeze()
        label_y = test_y.data.numpy()
        accuracy = sum(pred_y == label_y) / test_y.size()
        print("准确率为 %.2f" % (accuracy))

if __name__ == '__main__':
    eval()