def main(): args = options() torch.manual_seed(args.seed) torch.cuda.manual_seed_all(args.seed) np.random.seed(args.seed) testset = RegistrationData("PCRNet", ModelNet40Data(train=False, ), is_testing=True) test_loader = DataLoader( testset, batch_size=args.batch_size, shuffle=False, drop_last=False, num_workers=args.workers, ) if not torch.cuda.is_available(): args.device = "cpu" args.device = torch.device(args.device) # Create PointNet Model. ptnet = PointNet(emb_dims=args.emb_dims) model = iPCRNet(feature_model=ptnet) # Create sampler if args.sampler == "samplenet": sampler = SampleNet( args.num_out_points, args.bottleneck_size, args.projection_group_size, skip_projection=args.skip_projection, input_shape="bnc", output_shape="bnc", ) elif args.sampler == "fps": sampler = FPSSampler(args.num_out_points, permute=True, input_shape="bnc", output_shape="bnc") else: sampler = None model.sampler = sampler model = model.to(args.device) model = model.to(args.device) if args.pretrained: assert os.path.isfile(args.pretrained) model.load_state_dict(torch.load(args.pretrained, map_location="cpu")) model.to(args.device) test(args, model, test_loader)
def init_inference(): global device if args.model == 'resnet18': model = models.resnet18() model.fc = torch.nn.Linear(512, 3) elif args.model == 'samplenet': model = SampleNet() elif args.model == 'simplenet': model = SimpleNet() else: raise NotImplementedError() model.eval() model.load_state_dict(torch.load(args.pretrained_model)) model = model.cuda() x = torch.ones((1, 3, 240, 320)).cuda() from torch2trt import torch2trt #model_trt = torch2trt(model, [x], max_batch_size=100, fp16_mode=True) model_trt = torch2trt(model, [x], max_batch_size=100) torch.save(model_trt.state_dict(), args.trt_model)
def init_inference(): global model global device if args.model == 'resnet18': model = models.resnet18() model.fc = torch.nn.Linear(512, 3) elif args.model == 'samplenet': model = SampleNet() elif args.model == 'simplenet': model = SimpleNet() else: raise NotImplementedError() model.eval() #model.load_state_dict(torch.load(args.pretrained_model)) if args.trt_module: from torch2trt import TRTModule if args.trt_conversion: model.load_state_dict(torch.load(args.pretrained_model)) model = model.cuda() x = torch.ones((1, 3, 240, 320)).cuda() from torch2trt import torch2trt model_trt = torch2trt(model, [x], max_batch_size=100, fp16_mode=True) #model_trt = torch2trt(model, [x], max_batch_size=100) torch.save(model_trt.state_dict(), args.trt_model) exit() model_trt = TRTModule() #model_trt.load_state_dict(torch.load('road_following_model_trt_half.pth')) model_trt.load_state_dict(torch.load(args.trt_model)) model = model_trt.to(device) else: model.load_state_dict(torch.load(args.pretrained_model)) model = model.to(device)
def main(): # Parse arguments. args = parse_args() # Set device. device = 'cuda' if torch.cuda.is_available() else 'cpu' # Prepare dataset. np.random.seed(seed=0) image_dataframe = pd.read_csv(args.data_csv, engine='python', header=None) image_dataframe = image_dataframe.reindex( np.random.permutation(image_dataframe.index)) test_num = int(len(image_dataframe) * 0.2) train_dataframe = image_dataframe[test_num:] test_dataframe = image_dataframe[:test_num] train_data = MyDataset(train_dataframe, transform=transforms.ToTensor()) test_data = MyDataset(test_dataframe, transform=transforms.ToTensor()) train_loader = torch.utils.data.DataLoader(train_data, batch_size=20, shuffle=True) test_loader = torch.utils.data.DataLoader(test_data, batch_size=20) print('data set') # Set a model. if args.model == 'resnet18': model = models.resnet18() model.fc = torch.nn.Linear(512, 3) elif args.model == 'samplenet': model = SampleNet() elif args.model == 'simplenet': model = SimpleNet() else: raise NotImplementedError() model.train() model = model.to(device) print('model set') # Set loss function and optimization function. criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters(), lr=args.lr) #optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=0.9, weight_decay=5e-4) print('optimizer set') # Train and test. print('Train starts') for epoch in range(args.n_epoch): # Train and test a model. train_acc, train_loss = train(model, device, train_loader, criterion, optimizer) # Output score. if (epoch % args.test_interval == 0): test_acc, test_loss = test(model, device, test_loader, criterion) stdout_temp = 'epoch: {:>3}, train acc: {:<8}, train loss: {:<8}, test acc: {:<8}, test loss: {:<8}' print( stdout_temp.format(epoch + 1, train_acc, train_loss, test_acc, test_loss)) else: stdout_temp = 'epoch: {:>3}, train acc: {:<8}, train loss: {:<8}' #, test acc: {:<8}, test loss: {:<8}' print(stdout_temp.format(epoch + 1, train_acc, train_loss)) #, test_acc, test_loss)) # Save a model checkpoint. if (epoch % args.save_model_interval == 0 or epoch + 1 == args.n_epoch): model_ckpt_path = args.model_ckpt_path_temp.format( args.dataset_name, args.model_name, epoch + 1) torch.save(model.state_dict(), model_ckpt_path) print('Saved a model checkpoint at {}'.format(model_ckpt_path)) print('')
def main(): # Parse arguments. args = parse_args() # Set device. device = 'cuda' if torch.cuda.is_available() else 'cpu' ROOT_DIR = "" imgDataset = MyDataset(args.data_csv, ROOT_DIR, transform=transforms.ToTensor()) # Load dataset. train_data, test_data = train_test_split(imgDataset, test_size=0.2) pd.to_pickle(test_data, "test_data.pkl") del test_data train_loader = torch.utils.data.DataLoader(train_data, batch_size=20, shuffle=True) print('data set') # Set a model. if args.model == 'resnet18': model = models.resnet18() model.fc = torch.nn.Linear(512, 3) elif args.model == 'samplenet': model = SampleNet() else: raise NotImplementedError() model.train() model = model.to(device) print('model set') # Set loss function and optimization function. criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters(), lr=args.lr) #optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=0.9, weight_decay=5e-4) print('optimizer set') # Train and test. print('Train starts') for epoch in range(args.n_epoch): # Train and test a model. train_acc, train_loss = train(model, device, train_loader, criterion, optimizer) # Output score. if (epoch % args.test_interval == 0): pd.to_pickle(train_data, "train_data.pkl") del train_data test_data = pd.read_pickle("test_data.pkl") test_loader = torch.utils.data.DataLoader(test_data, batch_size=20, shuffle=True) del test_data test_acc, test_loss = test(model, device, test_loader, criterion) del test_loader stdout_temp = 'epoch: {:>3}, train acc: {:<8}, train loss: {:<8}, test acc: {:<8}, test loss: {:<8}' print( stdout_temp.format(epoch + 1, train_acc, train_loss, test_acc, test_loss)) train_data = pd.read_pickle("train_data.pkl") else: stdout_temp = 'epoch: {:>3}, train acc: {:<8}, train loss: {:<8}' #, test acc: {:<8}, test loss: {:<8}' print(stdout_temp.format(epoch + 1, train_acc, train_loss)) #, test_acc, test_loss)) # Save a model checkpoint. if (epoch % args.save_model_interval == 0 or epoch + 1 == args.n_epoch): model_ckpt_path = args.model_ckpt_path_temp.format( args.dataset_name, args.model_name, epoch + 1) torch.save(model.state_dict(), model_ckpt_path) print('Saved a model checkpoint at {}'.format(model_ckpt_path)) print('')
def main(): parser = sputils.get_parser() args = options(parser) torch.backends.cudnn.deterministic = True torch.manual_seed(args.seed) torch.cuda.manual_seed_all(args.seed) np.random.seed(args.seed) boardio = SummaryWriter(log_dir="checkpoints/" + args.exp_name) _init_(args) textio = IOStream("checkpoints/" + args.exp_name + "/run.log") textio.cprint(str(args)) # trainset = RegistrationData("PCRNet", ModelNet40Data(train=True, download=True)) # testset = RegistrationData("PCRNet", ModelNet40Data(train=False, download=True)) transforms = torchvision.transforms.Compose( [PointcloudToTensor(), OnUnitCube()]) traindata = ModelNetCls( args.num_in_points, transforms=transforms, train=True, download=False, folder=args.datafolder, ) testdata = ModelNetCls( args.num_in_points, transforms=transforms, train=False, download=False, folder=args.datafolder, ) trainset = RegistrationData("PCRNet", traindata) testset = RegistrationData("PCRNet", testdata) train_loader = DataLoader( trainset, batch_size=args.batch_size, shuffle=True, drop_last=True, num_workers=args.workers, ) test_loader = DataLoader( testset, batch_size=args.batch_size, shuffle=False, drop_last=False, num_workers=args.workers, ) if not torch.cuda.is_available(): args.device = "cpu" args.device = torch.device(args.device) # Create PointNet Model. ptnet = PointNet(emb_dims=args.emb_dims) model = iPCRNet(feature_model=ptnet) if args.train_pcrnet: model.requires_grad_(True) model.train() else: model.requires_grad_(False) model.eval() # Create sampler if args.sampler == "samplenet": sampler = SampleNet( args.num_out_points, args.bottleneck_size, args.projection_group_size, skip_projection=args.skip_projection, input_shape="bnc", output_shape="bnc", ) if args.train_samplenet: sampler.requires_grad_(True) sampler.train() else: sampler.requires_grad_(False) sampler.eval() else: sampler = None model.sampler = sampler model = model.to(args.device) checkpoint = None if args.resume: assert os.path.isfile(args.resume) checkpoint = torch.load(args.resume) args.start_epoch = checkpoint["epoch"] model.load_state_dict(checkpoint["model"]) if args.pretrained: assert os.path.isfile(args.pretrained) missing_keys, unexpected_keys = model.load_state_dict(torch.load( args.pretrained, map_location="cpu"), strict=False) if len(missing_keys) != 0: print(f"Found missing keys in checkpoint: {missing_keys}") if len(unexpected_keys) != 0: raise RuntimeError( f"Found missing keys in model: {unexpected_keys}") filtered_missing_keys = [ x for x in missing_keys if not x.startswith("sampler") ] if len(filtered_missing_keys) != 0: raise RuntimeError( f"Found missing keys in checkpoint: {filtered_missing_keys}") model.to(args.device) if args.eval: test(args, model, test_loader, textio) else: train(args, model, train_loader, test_loader, boardio, textio, checkpoint)