deep-learning - 将具有多个网络的 pytorch 模型转换为 onnx

Question

我正在尝试将具有多个网络的 pytorch 模型转换为 ONNX，但遇到了一些问题。

git 存储库:https://github.com/InterDigitalInc/HRFAE

训练师类:

class Trainer(nn.Module):
    def __init__(self, config):
        super(Trainer, self).__init__()
        # Load Hyperparameters
        self.config = config
        # Networks
        self.enc = Encoder()
        self.dec = Decoder()
        self.mlp_style = Mod_Net()
        self.dis = Dis_PatchGAN()
        ...

以下是经过训练的模型处理图像的方式:

    def gen_encode(self, x_a, age_a, age_b=0, training=False, target_age=0):
        if target_age:
            self.target_age = target_age
            age_modif = self.target_age*torch.ones(age_a.size()).type_as(age_a)
        else:
            age_modif = self.random_age(age_a, diff_val=25)

        # Generate modified image
        self.content_code_a, skip_1, skip_2 = self.enc(x_a)
        style_params_a = self.mlp_style(age_a)
        style_params_b = self.mlp_style(age_modif)
        
        x_a_recon = self.dec(self.content_code_a, style_params_a, skip_1, skip_2)
        x_a_modif = self.dec(self.content_code_a, style_params_b, skip_1, skip_2)
        
        return x_a_recon, x_a_modif, age_modif

下面是我如何转换为 onnx:

enc = Encoder()
dec = Decoder()
mlp = Mod_Net()
layers = [enc, mlp, dec]
model = torch.nn.Sequential(*layers)  
# here is my confusion: how do I specify the inputs of each layer?? 
# E.g. one of the outputs of 'enc' layer should be input of 'mlp' layer, 
# or the outputs of 'enc' layer should be part of inputs of 'dec' layer...

params = torch.load('./logs/001/checkpoint')  
model[0].load_state_dict(params['enc_state_dict'])
model[1].load_state_dict(params['mlp_style_state_dict'])
model[2].load_state_dict(params['dec_state_dict'])

torch.onnx.export(model, torch.randn([1, 3, 1024, 1024]), 'trained_hrfae.onnx', do_constant_folding=True)  

也许转换部分代码的方式有误？？谁能帮忙，非常感谢!

#20210629-11:52GMT 编辑:

我发现使用 torch.nn.Sequential 有限制。 Sequential 中前一层的输出应该与后一层的输入一致。所以我的代码根本不应该工作，因为“enc”层的输出与“mlp”层的输入不一致。

谁能帮忙把这种pytorch模型转换成onnx？非常感谢，再次:)

Answer 1

经过研究和尝试，我找到了一个可能正确的方法:

将每个网络(编码器、Mod_Net、解码器)转换为 onnx 模型，并在后面的逻辑过程或任何进一步的过程中处理它们的输入/输出(例如转换为 tflite 模型)。

我正在尝试使用此方法移植到 Android。

#Edit 20210705-03:52GMT#

另一种方法可能更好:写一个新的网络结合三个网络。我已经证明输出与原始 pytorch 模型相同。

class HRFAE(nn.Module):
    def __init__(self):
        super(HRFAE, self).__init__()
        self.enc = Encoder()
        self.mlp_style = Mod_Net()
        self.dec = Decoder()

    def forward(self, x, age_modif):
        content_code_a, skip_1, skip_2 = self.enc(x)
        style_params_b = self.mlp_style(age_modif)

        x_a_modif = self.dec(content_code_a, style_params_b, skip_1, skip_2)

        return x_a_modif

然后使用以下转换:

net = HRFAE()

params = torch.load('./logs/002/checkpoint')
net.enc.load_state_dict(params['enc_state_dict'])
net.mlp_style.load_state_dict(params['mlp_style_state_dict'])
net.dec.load_state_dict(params['dec_state_dict'])

net.eval()
torch.onnx.export(net, (torch.randn([1, 3, 512, 512]), torch.randn([1]).type(torch.long)), 'test_hrfae.onnx')

这应该是答案。