我正在尝试将pix2pix转换为可以在Lens Studio中运行的pb或onnx文件。Lens Studio对模型有严格的要求。我正在使用Lens Studio提供的这个指南来将这个PyTorch模型导出为onnx。问题在于这里找到的PyTorch模型使用了它自己的基类,而在示例中使用的是Module.nn,因此它没有torch.onnx.export
函数运行所需的方法/变量。到目前为止,我遇到了缺少一个名为training
的变量和一个名为train
的方法的问题。
尝试修改基础模型是否值得,还是应该尝试从头开始使用nn.Module构建它?有没有办法让pix2pix模型同时继承自抽象基类和nn.module?我是否没有理解情况?我之所以想按照Lens Studio的教程来做,是因为我已经通过不同的方式导出了onnx文件,但Lens Studio由于各种原因不接受这些文件。
另外,这是我第一次在SO上提问(在编程6年后),如果我犯了任何错误,请告诉我,我会更正。谢谢你。
这是来自教程的创建PyTorch模型用于Lens Studio的重要代码:
import torchimport torch.nn as nnclass Model(nn.Module): def __init__(self): super().__init__() self.layer = nn.Conv2d(in_channels=3, out_channels=1, kernel_size=3, stride=2, padding=1) def forward(self, x): out = self.layer(x) out = nn.functional.interpolate(out, scale_factor=2, mode='bilinear', align_corners=True) out = torch.nn.functional.softmax(out, dim=1) return out
我不会包含PyTorch模型的所有代码,因为它很大,但baseModel.py的开头是
import osimport torchfrom collections import OrderedDictfrom abc import ABC, abstractmethodfrom . import networksclass BaseModel(ABC): """这个类是一个模型的抽象基类(ABC)。 要创建子类,您需要实现以下五个函数: -- <__init__>: 初始化类;首先调用BaseModel.__init__(self, opt)。 -- <set_input>: 从数据集解包数据并应用预处理。 -- <forward>: 生成中间结果。 -- <optimize_parameters>: 计算损失、梯度并更新网络权重。 -- <modify_commandline_options>: (可选)添加模型特定的选项并设置默认选项。 """ def __init__(self, opt): """初始化BaseModel类。 参数: opt (Option类) -- 存储所有实验标志;需要是BaseOptions的子类 在创建自定义类时,您需要实现自己的初始化。 在此函数中,您应该首先调用<BaseModel.__init__(self, opt)> 然后,您需要定义四个列表: -- self.loss_names (字符串列表): 指定您想要绘制和保存的训练损失。 -- self.model_names (字符串列表): 定义我们在训练中使用的网络。 -- self.visual_names (字符串列表): 指定您想要显示和保存的图像。 -- self.optimizers (优化器列表): 定义并初始化优化器。您可以为每个网络定义一个优化器。如果两个网络同时更新,您可以使用itertools.chain来组合它们。参见cycle_gan_model.py的示例。 """ self.opt = opt self.gpu_ids = opt.gpu_ids self.isTrain = opt.isTrain self.device = torch.device('cuda:{}'.format(self.gpu_ids[0])) if self.gpu_ids else torch.device('cpu') # 获取设备名称:CPU或GPU self.save_dir = os.path.join(opt.checkpoints_dir, opt.name) # 将所有检查点保存到save_dir if opt.preprocess != 'scale_width': # 使用[scale_width],输入图像可能有不同的尺寸,这会损害cudnn.benchmark的性能。 torch.backends.cudnn.benchmark = True self.loss_names = [] self.model_names = [] self.visual_names = [] self.optimizers = [] self.image_paths = [] self.metric = 0 # 用于学习率策略'plateau'
对于pix2pix_model.py
import torchfrom .base_model import BaseModelfrom . import networksclass Pix2PixModel(BaseModel): """这个类实现了pix2pix模型,用于学习从输入图像到输出图像的映射,给定成对数据。 模型训练需要'--dataset_mode aligned'数据集。 默认情况下,它使用'--netG unet256'的U-Net生成器, '--netD basic'的判别器(PatchGAN), 和'--gan_mode'的香草GAN损失(在原始GAN论文中使用的交叉熵目标)。 pix2pix论文:https://arxiv.org/pdf/1611.07004.pdf """ @staticmethod def modify_commandline_options(parser, is_train=True): """添加新的数据集特定选项,并重写现有选项的默认值。 参数: parser -- 原始选项解析器 is_train (布尔值) -- 是否为训练阶段或测试阶段。您可以使用此标志来添加特定于训练或测试的选项。 返回: 修改后的解析器。 对于pix2pix,我们不使用图像缓冲区 训练目标是:GAN损失 + lambda_L1 * ||G(A)-B||_1 默认情况下,我们使用香草GAN损失、带批量规范化的UNet和对齐的数据集。 """ # 更改默认值以匹配pix2pix论文(https://phillipi.github.io/pix2pix/) parser.set_defaults(norm='batch', netG='unet_256', dataset_mode='aligned') if is_train: parser.set_defaults(pool_size=0, gan_mode='vanilla') parser.add_argument('--lambda_L1', type=float, default=100.0, help='L1损失的权重') return parser def __init__(self, opt): """初始化pix2pix类。 参数: opt (Option类) -- 存储所有实验标志;需要是BaseOptions的子类 """
(另外,如果你看到这篇文章并且看起来没有简单的出路,请告诉我,我知道刚开始某事时陷入太深的情况)
回答:
您的模型当然可以同时继承自基类和torch.nn.Module
(Python允许多重继承)。然而,您应该注意如果两个继承的类有相同名称的函数时可能发生的冲突(我至少看到了一个:它们的基类提供了eval
函数,而nn.module也是如此)。
但是,由于您不需要CycleGAN,并且很多代码是与他们的训练环境兼容的,您最好重新实现pix2pix。直接拿走代码,让它继承自nn.Module,从基类中复制粘贴有用/必需的函数,并将其全部转换为干净的PyTorch代码。您已经有了前向函数(这是PyTorch模块的唯一要求)。
他们使用的所有子网络(如resnet块)似乎已经继承自nn.Module
,所以这里不需要更改(但请仔细检查)。