PyTorch VAE无法转换为onnx

我正在尝试将PyTorch VAE转换为onnx，但我得到：

torch.onnx.symbolic.normal不存在

问题似乎源于

重新参数化（）

函数：

    def reparametrize(self, mu, logvar):
        std = logvar.mul(0.5).exp_()
        if self.have_cuda:
             eps = torch.normal(torch.zeros(std.size()),torch.ones(std.size())).cuda()
        else:
           eps = torch.normal(torch.zeros(std.size()),torch.ones(std.size()))
        return eps.mul(std).add_(mu)

我还尝试：

eps = torch.cuda.FloatTensor(std.size()).normal_()

产生错误的原因：

    Schema not found for node. File a bug report.
    Node: %173 : Float(1, 20) = aten::normal(%169, %170, %171, %172), scope: VAE 
    Input types:Float(1, 20), float, float, Generator

及

产生错误的原因：

    builtins.TypeError: i_(): incompatible function arguments. The following argument types are supported:
    1. (self: torch._C.Node, arg0: str, arg1: int) -> torch._C.Node
    Invoked with: %137 : Tensor = onnx::RandomNormal(), scope: VAE, 'shape', 133 defined in (%133 : int[] = prim::ListConstruct(%128, %132), scope: VAE) (occurred when translating randn)

我正在使用

cuda

任何想法都值得赞赏。也许我需要对onnx使用不同的

z

/潜伏期

---

注意：通过单步执行，我可以看到它正在为

torch.randn（）

查找

RandomNormal（）

，这应该是正确的。但是我在那一点上并没有访问参数的权限，所以我如何修复它呢？

简而言之，下面的代码可能会起作用。（至少在我的环境中，它在没有错误的情况下工作）

似乎

.size（）

运算符可能返回变量，而不是常量，因此它会导致onnx编译错误。（我在更改为使用.size（）时遇到了相同的错误）

---

太好了，谢谢！请注意，这是一个多么奇怪的问题，我现在从文档中看到，

torch.size（）

返回类型“torchSize”。。。所以，这让我通过了ONNX。。。现在CoreML正在抱怨RandonNormal。哈哈。。。多么惨败啊。

    builtins.TypeError: i_(): incompatible function arguments. The following argument types are supported:
    1. (self: torch._C.Node, arg0: str, arg1: int) -> torch._C.Node
    Invoked with: %137 : Tensor = onnx::RandomNormal(), scope: VAE, 'shape', 133 defined in (%133 : int[] = prim::ListConstruct(%128, %132), scope: VAE) (occurred when translating randn)

import torch
import torch.utils.data
from torch import nn
from torch.nn import functional as F



IN_DIMS = 28 * 28
BATCH_SIZE = 10
FEATURE_DIM = 20

class VAE(nn.Module):
    def __init__(self):
        super(VAE, self).__init__()

        self.fc1 = nn.Linear(784, 400)
        self.fc21 = nn.Linear(400, FEATURE_DIM)
        self.fc22 = nn.Linear(400, FEATURE_DIM)
        self.fc3 = nn.Linear(FEATURE_DIM, 400)
        self.fc4 = nn.Linear(400, 784)

    def encode(self, x):
        h1 = F.relu(self.fc1(x))
        return self.fc21(h1), self.fc22(h1)

    def reparameterize(self, mu, logvar):
        std = torch.exp(0.5*logvar)
        eps = torch.randn(BATCH_SIZE, FEATURE_DIM, device='cuda')
        return eps.mul(std).add_(mu)

    def decode(self, z):
        h3 = F.relu(self.fc3(z))
        return torch.sigmoid(self.fc4(h3))

    def forward(self, x):
        mu, logvar = self.encode(x)
        z = self.reparameterize(mu, logvar)
        recon_x = self.decode(z)

        return recon_x

model = VAE().cuda()

dummy_input = torch.randn(BATCH_SIZE, IN_DIMS, device='cuda')
torch.onnx.export(model, dummy_input, "vae.onnx", verbose=True)