|
- import torch
- import torch.nn as nn
- import os
- import math
- import logging
- import torch.utils.model_zoo as model_zoo
- import math
- import numpy as np
- import cv2
- from rknn.api import RKNN
- import rknn.api
- import torchvision.models as models
- model = nn.Sequential(
- nn.ConvTranspose2d(
- in_channels=1280,
- out_channels=16,
- kernel_size=4,
- stride=2,
- padding=1,
- output_padding=0,
- bias=True),
- # nn.BatchNorm2d(16, momentum=0.1),
- nn.ReLU(inplace=True),
- nn.ConvTranspose2d(
- in_channels=16,
- out_channels=16,
- kernel_size=4,
- stride=2,
- padding=1,
- output_padding=0,
- bias=True),
- nn.BatchNorm2d(16, momentum=0.1),
- nn.ReLU(inplace=True),
- nn.ConvTranspose2d(
- in_channels=16,
- out_channels=16,
- kernel_size=4,
- stride=2,
- padding=1,
- output_padding=0,
- bias=True),
- nn.BatchNorm2d(16, momentum=0.1),
- nn.ReLU(inplace=True),
- )
- trace_model = torch.jit.trace(model, torch.Tensor(1, 1280, 16, 16))
- trace_model.save('test_error.pt')
- input_size_list = [[1280, 16, 16]]
- # Create RKNN object
- rknn = RKNN(verbose=True)
- print('--> Loading model')
- ret = rknn.load_pytorch(model='test_error.pt', input_size_list=input_size_list)
第一个batchnorm如果注释掉,就能转成功;如果保留,就会出现上面的报错。请问这是为什么? |
|