Open 440981 opened 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 的名称(自己起名字) )
The code is as follows