def main(): # Parse command line arguments args = parse_args() # Create experiment summary = _make_experiment(args) # Create datasets train_data = KittiObjectDataset( args.root, 'train', args.grid_size, args.grid_res, args.yoffset) val_data = KittiObjectDataset( args.root, 'val', args.grid_size, args.grid_res, args.yoffset) # Apply data augmentation train_data = oft.AugmentedObjectDataset( train_data, args.train_image_size, args.train_grid_size, jitter=args.grid_jitter) # Create dataloaders train_loader = DataLoader(train_data, args.batch_size, shuffle=True, num_workers=args.workers, collate_fn=oft.utils.collate) val_loader = DataLoader(val_data, args.batch_size, shuffle=False, num_workers=args.workers,collate_fn=oft.utils.collate) # Build model model = OftNet(num_classes=1, frontend=args.frontend, topdown_layers=args.topdown, grid_res=args.grid_res, grid_height=args.grid_height) if len(args.gpu) > 0: torch.cuda.set_device(args.gpu[0]) model = nn.DataParallel(model, args.gpu).cuda() # Create encoder encoder = ObjectEncoder() # Setup optimizer optimizer = optim.SGD( model.parameters(), args.lr, args.momentum, args.weight_decay) scheduler = optim.lr_scheduler.ExponentialLR(optimizer, args.lr_decay) for epoch in range(1, args.epochs+1): print('\n=== Beginning epoch {} of {} ==='.format(epoch, args.epochs)) # Update and log learning rate scheduler.step(epoch-1) summary.add_scalar('lr', optimizer.param_groups[0]['lr'], epoch) # Train model train(args, train_loader, model, encoder, optimizer, summary, epoch) # Run validation every N epochs if epoch % args.val_interval == 0: validate(args, val_loader, model, encoder, summary, epoch) # Save model checkpoint save_checkpoint(args, epoch, model, optimizer, scheduler)
def main(): # Parse command line arguments args = parse_args() # Load validation dataset to visualise dataset = KittiObjectDataset( args.root, 'val', args.grid_size, args.grid_res, args.yoffset) # Build model model = OftNet(num_classes=1, frontend=args.frontend, topdown_layers=args.topdown, grid_res=args.grid_res, grid_height=args.grid_height) if args.gpu >= 0: torch.cuda.set_device(args.gpu) model.cuda() # Load checkpoint ckpt = torch.load(args.model_path) model = nn.DataParallel(model, [args.gpu]).cuda() model.load_state_dict(ckpt['model']) model = model.module # Create encoder encoder = ObjectEncoder(nms_thresh=args.nms_thresh) # Set up plots _, (ax1, ax2) = plt.subplots(nrows=2) plt.ion() # Iterate over validation images for _, image, calib, objects, grid in dataset: # Move tensors to gpu image = to_tensor(image) if args.gpu >= 0: image, calib, grid = image.cuda(), calib.cuda(), grid.cuda() # Run model forwards pred_encoded = model(image[None], calib[None], grid[None]) # Decode predictions pred_encoded = [t[0].cpu() for t in pred_encoded] detections = encoder.decode(*pred_encoded, grid.cpu()) image = image.cpu() calib = calib.cpu() # Visualize predictions visualize_objects(image, calib, detections, ax=ax1) ax1.set_title('Detections') visualize_objects(image, calib, objects, ax=ax2) ax2.set_title('Ground truth') plt.draw() plt.pause(0.01) plt.waitforbuttonpress()
def main(): # Parse command line arguments args = parse_args() # DLProf - Init PyProf if args.dlprof: pyprof.init(enable_function_stack=True) # Set num epochs to 1 if DLProf is enabled args.epochs = 1 # Create experiment summary = _make_experiment(args) # Create datasets train_data = KittiObjectDataset(args.root, 'train', args.grid_size, args.grid_res, args.yoffset) val_data = KittiObjectDataset(args.root, 'val', args.grid_size, args.grid_res, args.yoffset) # Apply data augmentation # train_data = oft.AugmentedObjectDataset( # train_data, args.train_image_size, args.train_grid_size, # jitter=args.grid_jitter) # Create dataloaders train_loader = DataLoader(train_data, args.batch_size, shuffle=True, num_workers=args.workers, collate_fn=oft.utils.collate) val_loader = DataLoader(val_data, args.batch_size, shuffle=False, num_workers=args.workers, collate_fn=oft.utils.collate) # Build model model = OftNet(num_classes=1, frontend=args.frontend, topdown_layers=args.topdown, grid_res=args.grid_res, grid_height=args.grid_height) if len(args.gpu) > 0: torch.cuda.set_device(args.gpu[0]) model = nn.DataParallel(model, args.gpu).cuda() # Create encoder encoder = ObjectEncoder() # Setup optimizer optimizer = optim.SGD(model.parameters(), args.lr, args.momentum, args.weight_decay) scheduler = optim.lr_scheduler.ExponentialLR(optimizer, args.lr_decay) # Creates a GradScaler once at the beginning of training for AMP. Created even if not being used. scaler = GradScaler() for epoch in range(1, args.epochs + 1): print('\n=== Beginning epoch {} of {} ==='.format(epoch, args.epochs)) # Update and log learning rate scheduler.step(epoch - 1) summary.add_scalar('lr', optimizer.param_groups[0]['lr'], epoch) # Train model if args.dlprof: with torch.autograd.profiler.emit_nvtx(): train(args, train_loader, model, encoder, optimizer, summary, epoch, scaler) else: train(args, train_loader, model, encoder, optimizer, summary, epoch, scaler) # Run validation every N epochs if epoch % args.val_interval == 0: if args.dlprof: with torch.autograd.profiler.emit_nvtx(): validate(args, val_loader, model, encoder, summary, epoch) else: validate(args, val_loader, model, encoder, summary, epoch) # Save model checkpoint save_checkpoint(args, epoch, model, optimizer, scheduler)