转换为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