onnx export&tensorrt engine build
SoulProficiency opened this issue · comments
感谢NeRCo的贡献，目前想将NeRCo部署在条件较为苛刻的现实场景。
请问官方可以给出预训练模型的onnx格式导出吗？
谢谢 :)
您也可以通过PyTorch提供的torch.onnx.export函数转化为ONNX格式
谢谢 :) 您也可以通过PyTorch提供的torch.onnx.export函数转化为ONNX格式
目前依据test.py作为模板进行onnx导出的操作，以下是我的代码：
model = create_model(opt) # create a model given opt.model and other options
model.setup(opt) # regular setup: load and print networks; create schedulers
model.print_networks(True)
model.eval()
x = torch.randn(1, 3, 640, 640, requires_grad=True)
torch.onnx.export(model, x, "enhance.onnx", opset_version=17)
模型可以打印，但有如下错误：
Traceback (most recent call last):
File "./test.py", line 23, in
torch.onnx.export(model, x, "enhance.onnx", opset_version=17)
File "x\lib\site-packages\torch\onnx\utils.py", line 506, in export
_export(
File "x\lib\site-packages\torch\onnx\utils.py", line 1525, in _export
with exporter_context(model, training, verbose):
File "x\lib\contextlib.py", line 113, in enter
return next(self.gen)
File "x\lib\site-packages\torch\onnx\utils.py", line 178, in exporter_context
with select_model_mode_for_export(
File "x\lib\contextlib.py", line 113, in enter
return next(self.gen)
File "x\lib\site-packages\torch\onnx\utils.py", line 139, in disable_apex_o2_state_dict_hook
for module in model.modules():
AttributeError: 'NeRComodel' object has no attribute 'modules'
本机的torch版本：'2.0.0+cu118'
你可以尝试不直接把NeRCo作为整体导出，而是分别加载并导出3个子模型（latest_net_G_A.pth，latest_net_H.pth，latest_net_Pre.pth）。
你可以尝试不直接把NeRCo作为整体导出，而是分别加载并导出3个子模型（latest_net_G_A.pth，latest_net_H.pth，latest_net_Pre.pth）。
好的，感谢回复
如果有什么问题欢迎提问 ：）我会尽量解答
如果有什么问题欢迎提问 ：）我会尽量解答
作者您好，我大致阅读了文章，并做了一定总结，如果目前仅需要对图像进行增强，应该不需要像训练阶段同时使用两张图片吧(paper section3.1 paragraphs 3)。结合您的提示，我目前个人的想法如下：
1.直接使用Enhance Module对图像进行增强(paper Figure 2 top-right)。这种做法应该直接导出latest_net_H.pth(如果我没猜错的话)就可以了。
2.借助Mask Extractor和Neural Representation Normalization作为图像预处理的手段，然后通过Enhance Module进行增强(paper figure2 top-left)。这样仍需要导出Mask Extractor和Neural Representation Normalization的onnx文件。
以下是我个人的一些问题：
1.直接运行test.py，会提示文件的下载(默认应该是resnet50)，这应该是encoder权重吧(至少大多数论文是这么做的)。
2.上述提及的两种方式是否理解正确？如果存在问题请作者纠正。
3.模型分开导出的话，latest_net_H,latest_net_Pre,latest_net_G_A分别对应哪些部分，Mask Extractor和Neural Representation Normalization的权重又在哪些权重文件里面？
1、latest_net_Pre对应NRN；latest_net_H对应Mask Extractor；latest_net_G_A对应Enhance Module。
2、所以你上面理解的第一条有误，如果你只用Enhance Module应该下载latest_net_G_A。但我不建议你这么做，因为latest_net_G_A的输入通道是6（注意NRN最后输出的是concat结果的6通道图片），你只用这个预训练模型可能会有通道维度的问题。
3、目前我猜测你那边运行test.py有文件下载可能是要下载CLIP模型，你可以提供你的运行界面来进一步展示具体情况。
1.首次运行下载的内容如下：
2.参考代码结合您的描述，正确的方式依旧会使用到nrn，maskextractor和enhance module。我参考了模型的forward代码，加上了自己的注释(以推理的视角进行阐述，假定self.real_A为待推理的单张图片)：
def forward(self):
# nrn output
self.pre_A = self.netPre(self.real_A)
# mask extractor output
self.H, self.mask = self.netH(self.real_A)
# 通道合并(之前您提及的)
temp = torch.cat((self.real_A, self.pre_A), 1)
# mask和合并后的结果相乘并通过enhance module进行增强
# 此处fake_B即为增强后的结果(对于前向推理而言，到此已经可以结束)
self.fake_B = self.netG_A(temp * self.mask) # G_A(A)
self.rec_A = self.netG_B(self.fake_B) # G_B(G_A(A))
self.fake_A = self.netG_B(self.real_B * self.mask) # G_B(B)
self.pre_A1 = self.netPre(self.fake_A)
temp = torch.cat((self.fake_A, self.pre_A1), 1)
self.rec_B = self.netG_A(temp) # G_A(G_B(B))
不知上述的理解是否正确？netG_B是训练的Degrade Module用于和enhance Module进行对抗训练。
理解是对的，初次下载的权重应该就是CLIP的权重
理解是对的，初次下载的权重应该就是CLIP的权重
感谢~
欢迎随时提问：）我暂时关闭这个issue，如果你需要可以随时打开它。