|
7#
楼主 |
发表于 2020-4-13 17:47:46
|
只看该作者
后续进行模型转换的时候又遇到了permute转换失效和core dumped两个错误。
模型定义:
import torch
import torch.nn as nn
class ClassHead(nn.Module):
def __init__(self, inchannels=512, num_anchors=3):
super(ClassHead, self).__init__()
self.num_anchors = num_anchors
self.conv1x1 = nn.Conv2d(inchannels, self.num_anchors * 2, kernel_size=(1, 1), stride=1, padding=0)
self.output_act = nn.LogSoftmax(dim=-1)
def forward(self, x):
out = self.conv1x1(x)
out = out.permute(0, 2, 3, 1)
#b, h, w, c = out.shape
#out = out.view(b, h, w, self.num_anchors, 2)
#out = self.output_act(out)
return out
classhead = ClassHead(512,3)
out = classhead(torch.Tensor(1,512,28,28))
print(out.shape)
trace_model = torch.jit.trace(classhead, torch.Tensor(1,512,28,28))
trace_model.save('./pt_out/classhead_output.pt')
rknn转化和推理:
import numpy as np
import cv2
from rknn.api import RKNN
import torchvision.models as models
import torch
if __name__ == '__main__':
model = './pt_out/classhead_output.pt'
input_size_list = [[512,28,28]]
# Create RKNN object
rknn = RKNN(verbose=True)
# Load pytorch model
print('--> Loading model')
ret = rknn.load_pytorch(model=model, input_size_list=input_size_list)
if ret != 0:
print('Load pytorch model failed!')
exit(ret)
print('done')
# Build model
print('--> Building model')
ret = rknn.build(do_quantization=True, dataset = './dataset.txt')
if ret != 0:
print('Build pytorch failed!')
exit(ret)
print('done')
# Export rknn model
print('--> Export RKNN model')
ret = rknn.export_rknn('./pt_out/classhead_output.rknn')
if ret != 0:
print('Export rknn failed!')
exit(ret)
print('done')
ret = rknn.load_rknn('pt_out/classhead_output.rknn')
# init runtime environment
print('--> Init runtime environment')
ret = rknn.init_runtime()
if ret != 0:
print('Init runtime environment failed')
exit(ret)
print('done')
# Set inputs
img = np.load('./input.npy')
# Inference
print('--> Running model')
outputs = rknn.inference(inputs=[img])
print('output')
print(outputs[0].shape)
在这里维度转换out = out.permute(0, 2, 3, 1)会失效。torch模型推理输出尺寸为(1, 28, 28, 6),但是rknn模型推理的输出尺寸则为(1, 6, 28, 28),permute函数的维度转换没有生效。
第二个问题是如果如下所示继续加入view和Logsoftmax后:
import torch
import torch.nn as nn
class ClassHead(nn.Module):
def __init__(self, inchannels=512, num_anchors=3):
super(ClassHead, self).__init__()
self.num_anchors = num_anchors
self.conv1x1 = nn.Conv2d(inchannels, self.num_anchors * 2, kernel_size=(1, 1), stride=1, padding=0)
self.output_act = nn.LogSoftmax(dim=-1)
def forward(self, x):
out = self.conv1x1(x)
out = out.permute(0, 2, 3, 1)
b, h, w, c = out.shape
out = out.view(b, h, w, self.num_anchors, 2)
out = self.output_act(out)
return out
classhead = ClassHead(512,3)
out = classhead(torch.Tensor(1,512,28,28))
print(out.shape)
trace_model = torch.jit.trace(classhead, torch.Tensor(1,512,28,28))
trace_model.save('./pt_out/classhead_output_2.pt')
模型转换过程会在build时发出warning: Do not support shape > 4. 虽然能够转换生成rknn模型但是会在Init runtime environment时报错:Segmentation fault (core dumped)
我在dataset中指定的input.npy是一个shape为(1, 28, 28, 512)的一个numpy数组。请问这两个错误是什么原因
代码以及模型、错误记录:
链接:https://pan.baidu.com/s/1FlaC6CSMhRzaewiX2iJnsQ
提取码:ovnh
|
|