PyTorch To ONNX失败,出现以下错误

时间:2019-12-12 15:19:55

标签: pytorch onnx

转换为ONNX时弹出以下错误

我正在尝试将VAE模型转换为ONNX格式,但是下面存在一些问题

TypeError:i_():不兼容的函数参数。支持以下参数类型:     1.(self:torch._C.Node,arg0:str,arg1:int)-> torch._C.Node

请在下面找到要转换的代码

generator = VAEGen(3).cuda() # Initalize your weights there
dummy_input = torch.randn(1 ,3, 256, 256).cuda()
torch.onnx.export(generator, dummy_input, "exported.onnx")

并在下面找到VAE代码

class VAEGen(nn.Module):
# VAE architecture
def __init__(self, input_dim):
    super(VAEGen, self).__init__()
    # dim = params['dim']
    # n_downsample = params['n_downsample']
    # n_res = params['n_res']
    # activ = params['activ']
    # pad_type = params['pad_type']

    # content encoder
    self.enc = ContentEncoder(2, 4, 3, 64, 'in', 'relu', pad_type='reflect')
    self.dec = Decoder(2, 4, self.enc.output_dim, 3, res_norm='in', activ='relu', pad_type='reflect')

def forward(self, images):
    # This is a reduced VAE implementation where we assume the outputs are multivariate Gaussian distribution with mean = hiddens and std_dev = all ones.
    hiddens,_ = self.encode(images)
    if self.training == False:
        noise = Variable(torch.randn(hiddens.size()).cuda(hiddens.data.get_device()))
        images_recon = self.decode(hiddens + noise)
    else:
        images_recon = self.decode(hiddens)
    return images_recon, hiddens

def encode(self, images):
    hiddens = self.enc(images)
    noise = Variable(torch.randn(hiddens.size()).cuda(hiddens.data.get_device()))
    return hiddens, noise

def decode(self, hiddens):
    images = self.dec(hiddens)
    return images

0 个答案:

没有答案