Only tensors or tuples of tensors can be output from traced functions错误解决

(TorchScript应用) PyTorch模型转换为Torch脚本的代码出错。

出现原因：想将pytorch训练的.pth文件转成C++能处理的.pt文件。用的TorchScript的方法。具体代码如下：

import argparse

import cv2

import torchvision

import numpy as np

import torch

import torch.nn.functional as F

from torchvision import transforms

from PIL import Image

from tqdm import tqdm

from unet import NestedUNet

from unet import UNet

from utils.dataset import BasicDataset

from config import UNetConfig

cfg = UNetConfig()

device = torch.device('cpu')

model = eval(cfg.model)(cfg)

path = 'data/checkpoints/epoch_9.pth'#自己模型训练的结果

model.load_state_dict(torch.load(path, map_location=device))

model.to(device=device)

print(model)img_path = 'data/00000.jpg'#自己模型准备用的数据

img = cv2.imread(img_path)

imgXX = Image.fromarray(cv2.cvtColor(img,cv2.COLOR_BGR2RGB))

img1 = torch.from_numpy(BasicDataset.preprocess(imgXX, cfg.scale))

img1 = img1.unsqueeze(0)

img1 = img1.to(device=device, dtype=torch.float32)

traced_script_module = torch.jit.trace(model,img1)#报错的代码位置

traced_script_module.save("torch_script_eval.pt")

报错的内容为：Only tensors or tuples of tensors can be output from traced functions

返回的是字典，是不支持的。想办法把返回的字典变成 tensors，就可以了。

怀疑model这个模型的返回值非法了。找到model模型定义的返回值位置

def forward(self, input):

x0_0 = self.conv0_0(input)

x1_0 = self.conv1_0(self.pool(x0_0))

x0_1 = self.conv0_1(torch.cat([x0_0, self.up(x1_0)], 1))

x2_0 = self.conv2_0(self.pool(x1_0))

x1_1 = self.conv1_1(torch.cat([x1_0, self.up(x2_0)], 1))

x0_2 = self.conv0_2(torch.cat([x0_0, x0_1, self.up(x1_1)], 1))

x3_0 = self.conv3_0(self.pool(x2_0))

x2_1 = self.conv2_1(torch.cat([x2_0, self.up(x3_0)], 1))

x1_2 = self.conv1_2(torch.cat([x1_0, x1_1, self.up(x2_1)], 1))

x0_3 = self.conv0_3(torch.cat([x0_0, x0_1, x0_2, self.up(x1_2)], 1))

x4_0 = self.conv4_0(self.pool(x3_0))

x3_1 = self.conv3_1(torch.cat([x3_0, self.up(x4_0)], 1))

x2_2 = self.conv2_2(torch.cat([x2_0, x2_1, self.up(x3_1)], 1))

x1_3 = self.conv1_3(torch.cat([x1_0, x1_1, x1_2, self.up(x2_2)], 1))

x0_4 = self.conv0_4(torch.cat([x0_0, x0_1, x0_2, x0_3, self.up(x1_3)], 1))

if self.deepsupervision:

output1 = self.final1(x0_1)

output2 = self.final2(x0_2)

output3 = self.final3(x0_3)

output4 = self.final4(x0_4)

return [output1, output2, output3, output4)] #按程序设定 走的是这个分支。 显然[ ]不符合。

else:

output = self.final(x0_4)

return output

将其改为return （output1, output2, output3, output4)）#报错解决了。

总结：raced_script_module = torch.jit.trace(model,img1)这个位置报这类错，就怀疑返回值问题，然后一步步查。

上面这个改可能不符合代码原意了。再看看