zju3dv / NeuralRecon

Code for "NeuralRecon: Real-Time Coherent 3D Reconstruction from Monocular Video", CVPR 2021 oral
https://zju3dv.github.io/neuralrecon/
Apache License 2.0
2.03k stars 294 forks source link

Export the ONNX newspaper too many indices for tensor of dimension 0 #137

Open 440981 opened 1 year ago

440981 commented 1 year ago

The code is as follows

import torch
from torchvision import models
from torch.utils.data import DataLoader

# 有 GPU 就用 GPU,没有就用 CPU
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print('device', device)
import argparse
from models import NeuralRecon
from config import cfg, update_config

parser = argparse.ArgumentParser(description='NeuralRecon Real-time Demo')
parser.add_argument('--cfg',
                    help='experiment configure file name',
                    default='/mnt/workspace/NeuralRecon/config/demo.yaml',
                    #required=True,
                    type=str)

parser.add_argument('opts',
                    help="Modify config options using the command-line",
                    default=None,
                    nargs=argparse.REMAINDER)

# parse arguments and check
args = parser.parse_args(args=[])
update_config(cfg, args)
model = NeuralRecon(cfg).cuda().eval()
model = torch.nn.DataParallel(model, device_ids=[0])
state_dict = torch.load('/mnt/workspace/NeuralRecon/checkpoint/model_000047.ckpt')
model.load_state_dict(state_dict['model'], strict=False)
from datasets import find_dataset_def, transforms
transform = [transforms.ResizeImage((640, 480)),
             transforms.ToTensor(),
             transforms.RandomTransformSpace(
                 [96, 96, 96], 0.04, random_rotation=False, random_translation=False,
                 paddingXY=0, paddingZ=0, max_epoch=991),
             transforms.IntrinsicsPoseToProjection(9, 9)]

transforms = transforms.Compose(transform)
ARKitDataset = find_dataset_def('demo')
test_dataset = ARKitDataset('/mnt/workspace/NeuralRecon/data/2023-08-25T08-43-54', "test", transforms, 9, len([0, 0, 0]) - 1)
data_loader = DataLoader(test_dataset, 1, shuffle=False, num_workers=1, drop_last=False)
#for frag_idx, sample in enumerate(data_loader):
#    print(sample)
enum_fruits = list(enumerate(data_loader))
frag_idx, sample = enum_fruits[0]
frag_len = len(data_loader)
save_scene = (True and frag_idx == frag_len - 1) or False or False
frag_idx, sample = enum_fruits[0]
outputs, loss_dict = model(sample, save_scene)

with torch.no_grad():
    torch.onnx.export(
        model.module,                   # 要转换的模型
        sample,save_scene,                       # 模型的任意一组输入
        '/mnt/workspace/NeuralRecon/NeuralRecon.onnx', # 导出的 ONNX 文件名
        opset_version=11,        # ONNX 算子集版本
        input_names=['sample','save_scene'],   # 输入 Tensor 的名称(自己起名字)
        output_names=['outputs','loss_dict']  # 输出 Tensor 的名称(自己起名字)
    ) 
image