NeuralRecon
NeuralRecon copied to clipboard
Export the ONNX newspaper too many indices for tensor of dimension 0
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 的名称(自己起名字)
)