PyTorch - 将 ProGAN 代理从 pth 转换为 onnx

Question

我使用此 PyTorch 重新实现训练了一个 ProGAN 代理，并将该代理另存为 .现在我需要将代理转换为格式，我正在使用此scipt执行此操作：.pth.onnx

from torch.autograd import Variable

import torch.onnx

import torchvision

import torch

device = torch.device("cuda")

dummy_input = torch.randn(1, 3, 64, 64)

state_dict = torch.load("GAN_agent.pth", map_location = device)

torch.onnx.export(state_dict, dummy_input, "GAN_agent.onnx")

一旦我运行它，我得到错误（下面的完整提示）。据我所知，问题在于将代理转换为.onnx需要更多信息。我错过了什么吗？AttributeError: 'collections.OrderedDict' object has no attribute 'state_dict'

---------------------------------------------------------------------------

AttributeError                            Traceback (most recent call last)

<ipython-input-2-c64481d4eddd> in <module>

     10 state_dict = torch.load("GAN_agent.pth", map_location = device)

     11 

---> 12 torch.onnx.export(state_dict, dummy_input, "GAN_agent.onnx")

~\anaconda3\envs\Basemap_upres\lib\site-packages\torch\onnx\__init__.py in export(model, args, f, export_params, verbose, training, input_names, output_names, aten, export_raw_ir, operator_export_type, opset_version, _retain_param_name, do_constant_folding, example_outputs, strip_doc_string, dynamic_axes, keep_initializers_as_inputs)

    146                         operator_export_type, opset_version, _retain_param_name,

    147                         do_constant_folding, example_outputs,

--> 148                         strip_doc_string, dynamic_axes, keep_initializers_as_inputs)

    149 

    150 

Answer 1

您拥有的文件是 ，它们只是图层名称到权重偏差和类似值的映射（有关更全面的介绍，请参阅此处）。state_dicttensor这意味着你需要一个模型，以便可以映射那些节省的权重和偏差，但首先要做的事情是：1. 模型准备克隆模型定义所在的存储库并打开文件 。我们需要进行一些修改才能使其与 . 导出器需要仅作为（或/个）传递，而类需要和参数）。/pro_gan_pytorch/pro_gan_pytorch/PRO_GAN.pyonnxonnxinputtorch.tensorlistdictGeneratorintfloat简单的解决方案是稍微修改一下函数（文件中的行，可以在GitHub上验证它）到以下内容：forward80def forward(self, x, depth, alpha):&nbsp; &nbsp; """&nbsp; &nbsp; forward pass of the Generator&nbsp; &nbsp; :param x: input noise&nbsp; &nbsp; :param depth: current depth from where output is required&nbsp; &nbsp; :param alpha: value of alpha for fade-in effect&nbsp; &nbsp; :return: y => output&nbsp; &nbsp; """&nbsp; &nbsp; # THOSE TWO LINES WERE ADDED&nbsp; &nbsp; # We will pas tensors but unpack them here to `int` and `float`&nbsp; &nbsp; depth = depth.item()&nbsp; &nbsp; alpha = alpha.item()&nbsp; &nbsp; # THOSE TWO LINES WERE ADDED&nbsp; &nbsp; assert depth < self.depth, "Requested output depth cannot be produced"&nbsp; &nbsp; y = self.initial_block(x)&nbsp; &nbsp; if depth > 0:&nbsp; &nbsp; &nbsp; &nbsp; for block in self.layers[: depth - 1]:&nbsp; &nbsp; &nbsp; &nbsp; &nbsp; &nbsp; y = block(y)&nbsp; &nbsp; &nbsp; &nbsp; residual = self.rgb_converters[depth - 1](self.temporaryUpsampler(y))&nbsp; &nbsp; &nbsp; &nbsp; straight = self.rgb_converters[depth](self.layers[depth - 1](y))&nbsp; &nbsp; &nbsp; &nbsp; out = (alpha * straight) + ((1 - alpha) * residual)&nbsp; &nbsp; else:&nbsp; &nbsp; &nbsp; &nbsp; out = self.rgb_converters[0](y)&nbsp; &nbsp; return out此处仅添加了解包方式。每个不属于类型的输入都应在函数定义中打包为一个，并在函数顶部尽快解压缩。它不会破坏您创建的检查点，所以不用担心，因为它只是映射。item()Tensorlayer-weight2. 模型导出将此脚本放在 （位置也位于）：/pro_gan_pytorchREADME.mdimport torchfrom pro_gan_pytorch import PRO_GAN as pggen = torch.nn.DataParallel(pg.Generator(depth=9))gen.load_state_dict(torch.load("GAN_GEN_SHADOW_8.pth"))module = gen.module.to("cpu")# Arguments like depth and alpha may need to be changeddummy_inputs = (torch.randn(1, 512), torch.tensor([5]), torch.tensor([0.1]))torch.onnx.export(module, dummy_inputs, "GAN_GEN8.onnx", verbose=True)请注意以下几点：我们必须在加载权重之前创建模型，因为它是唯一的。state_dicttorch.nn.DataParallel是必需的，因为这是模型的训练对象（不确定您的情况，请相应地进行调整）。加载后，我们可以通过属性获取模块本身。module一切都被扔到了，我认为没有必要在这里。如果你坚持的话，你可以把一切都扔到。CPUGPUGPU生成器的虚拟输入不能是图像（我使用了存储库作者在其Google云端硬盘上提供的文件），它必须是带有元素的噪音。512运行它，你的文件应该在那里。.onnx哦，由于您遵循不同的检查点，您可能希望遵循类似的过程，尽管不能保证一切都会正常工作（尽管它看起来确实如此）。