1.环境
- Pytorch1.2 cuda 10.0
- Windows 10
- Anaconda
- libtorch1.2 release cuda version
- VS2017
2. 描述
对比Tensortflow框架，pytorch在深度学习模型的研究和产品化方面做了一个快速转换的通道，将算法设计人员和程序开发人员所擅长的东西无缝链接起来。笔者用标准的Unet网络进行图像语义分割部署实验。
2.1 Unet Model
class UNetScript(nn.Module):
def __init__(self,in_ch,out_ch):
super(UNetScript,self).__init__()
self.conv1 = DoubleConv(in_ch,64)
self.pool1 = nn.MaxPool2d(2)#resize the image half each time
self.conv2 = DoubleConv(64,128)
self.pool2 = nn.MaxPool2d(2)
self.conv3 = DoubleConv(128,256)
self.pool3 = nn.MaxPool2d(2)
self.conv4 = DoubleConv(256,512)
self.pool4 = nn.MaxPool2d(2)
self.conv5 = DoubleConv(512,1024)
#invert convolution
self.up6 = nn.ConvTranspose2d(1024,512,2,stride=2)
self.conv6 = DoubleConv(1024,512)
self.up7 = nn.ConvTranspose2d(512,256,2,stride=2)
self.conv7 = DoubleConv(512,256)
self.up8 = nn.ConvTranspose2d(256,128,2,stride=2)
self.conv8 = DoubleConv(256,128)
self.up9 = nn.ConvTranspose2d(128,64,2,stride=2)
self.conv9 = DoubleConv(128,64)
self.conv10 = nn.Conv2d(64,out_ch,1)
self.sigmoid = nn.Sigmoid()
def forward(self,x):
c1 = self.conv1(x)
p1 = self.pool1(c1)
c2 = self.conv2(p1)
p2 = self.pool2(c2)
c3 = self.conv3(p2)
p3 = self.pool3(c3)
c4 = self.conv4(p3)
p4 = self.pool4(c4)
c5 = self.conv5(p4)
up_6 = self.up6(c5)
merge6 = torch.cat([up_6,c4],dim=1)#splice according to dimension 1(colum), colum increase
c6 = self.conv6(merge6)
up_7 = self.up7(c6)
merge7 = torch.cat([up_7,c3],dim=1)
c7 = self.conv7(merge7)
up_8 = self.up8(c7)
merge8 = torch.cat([up_8,c2],dim=1)
c8 = self.conv8(merge8)
up_9 = self.up9(c8)
merge9 = torch.cat([up_9,c1],dim=1)
c9 = self.conv9(merge9)
c10 = self.conv10(c9)
if x.dim() > 1:
out = self.sigmoid(c10)#normalize to[0, 1]
else:
out = self.sigmoid(c10)#normalize to[0, 1]
out = out.mul(255)
return out2.2 Main Function
笔者已经提前用上面的模型训练出了参数（weights_0.pth）
#method 1
checkpointUrl = "./pytorchUnet/weights_0.pth"
network = UNetScript(in_ch = 3, out_ch = 1)
stateDict = torch.load(checkpointUrl)
network.load_state_dict(stateDict)
scriptModel = torch.jit.script(network)
scriptModel.save("unet_script_model_1.pt")
#method 2
unet_model = UNetScript(in_ch = 3, out_ch = 1)
unet_script = torch.jit.script(unet_model)
unet_script.save("unet_script_model_2.pt")2.3 坑都有哪些
2.3.1 trace和script功能使用场景有区别
对比torch.jit.trace功能，此处foward函数里面有依赖于输入数据的控制流，故必须使用torch.jit.script功能。
2.3.2 模型中的注释必须是全英文
有些人说，中文不可以，日文可以不，也不行，笔者当时就是看到一篇日文帖子才发现这个坑，日文好的同学可以去看看下面这个链接，不好的也可以看看，有道做助手。
https://qiita.com/Naomi310/items/1a60bb8d374c41d490e6
issue as follows:
2.3.3 模型里面要严格遵守函数在init初始化，forward来使用
看下面这个例子：
out = nn.Sigmoid()(c10)#normalize to[0, 1]我在forward里面使用未经init初始化的nn.Simoid，在script module时出现以下错误：
这个错误不容易发现，报的错误提示与实际解决方案隔着一个川普，因为笔者使用trace功能在libtorch测试下完全OK.所以...程序员还是有个好的编程习惯，就因为这个，花了一天时间把它解决。
扫描二维码关注公众号，回复：
12434791 查看本文章
2.3.4 训练模型和序列化模型有区别
什么意思呢，还是以上面的Unet模型为例，如果序列化时对原来forward函数做了无关紧要的修改（不影响函数输出），这时测试时出来的结果是不对的，但也不会报错。慎之。
2.3.5 序列化模型三种方法
- 训练之后，马上调用script函数进行模型序列化。
- 加载已经训练过的参数（weights），然后采用main函数中的method 1进行序列化。
- 按照mian函数中的method2进行序列化
经过笔者测试，直接给答案，前面两种OK，第三种有问题，测试结果不对，但程序也不会报错。
因为图片涉密，此处没有show the result，但应该有掌声。