Spaces:
Running
Running
| import argparse | |
| import numpy as np | |
| import os | |
| import shutil | |
| import torch | |
| import torch.optim as optim | |
| from torch.utils.data import DataLoader | |
| from tqdm import tqdm | |
| import warnings | |
| from lib.dataset import MegaDepthDataset | |
| from lib.exceptions import NoGradientError | |
| from lib.loss import loss_function | |
| from lib.model import D2Net | |
| # CUDA | |
| use_cuda = torch.cuda.is_available() | |
| device = torch.device("cuda:0" if use_cuda else "cpu") | |
| # Seed | |
| torch.manual_seed(1) | |
| if use_cuda: | |
| torch.cuda.manual_seed(1) | |
| np.random.seed(1) | |
| # Argument parsing | |
| parser = argparse.ArgumentParser(description="Training script") | |
| parser.add_argument( | |
| "--dataset_path", type=str, required=True, help="path to the dataset" | |
| ) | |
| parser.add_argument( | |
| "--scene_info_path", type=str, required=True, help="path to the processed scenes" | |
| ) | |
| parser.add_argument( | |
| "--preprocessing", | |
| type=str, | |
| default="caffe", | |
| help="image preprocessing (caffe or torch)", | |
| ) | |
| parser.add_argument( | |
| "--model_file", type=str, default="models/d2_ots.pth", help="path to the full model" | |
| ) | |
| parser.add_argument( | |
| "--num_epochs", type=int, default=10, help="number of training epochs" | |
| ) | |
| parser.add_argument("--lr", type=float, default=1e-3, help="initial learning rate") | |
| parser.add_argument("--batch_size", type=int, default=1, help="batch size") | |
| parser.add_argument( | |
| "--num_workers", type=int, default=4, help="number of workers for data loading" | |
| ) | |
| parser.add_argument( | |
| "--use_validation", | |
| dest="use_validation", | |
| action="store_true", | |
| help="use the validation split", | |
| ) | |
| parser.set_defaults(use_validation=False) | |
| parser.add_argument( | |
| "--log_interval", type=int, default=250, help="loss logging interval" | |
| ) | |
| parser.add_argument("--log_file", type=str, default="log.txt", help="loss logging file") | |
| parser.add_argument( | |
| "--plot", dest="plot", action="store_true", help="plot training pairs" | |
| ) | |
| parser.set_defaults(plot=False) | |
| parser.add_argument( | |
| "--checkpoint_directory", | |
| type=str, | |
| default="checkpoints", | |
| help="directory for training checkpoints", | |
| ) | |
| parser.add_argument( | |
| "--checkpoint_prefix", | |
| type=str, | |
| default="d2", | |
| help="prefix for training checkpoints", | |
| ) | |
| args = parser.parse_args() | |
| print(args) | |
| # Create the folders for plotting if need be | |
| if args.plot: | |
| plot_path = "train_vis" | |
| if os.path.isdir(plot_path): | |
| print("[Warning] Plotting directory already exists.") | |
| else: | |
| os.mkdir(plot_path) | |
| # Creating CNN model | |
| model = D2Net(model_file=args.model_file, use_cuda=use_cuda) | |
| # Optimizer | |
| optimizer = optim.Adam( | |
| filter(lambda p: p.requires_grad, model.parameters()), lr=args.lr | |
| ) | |
| # Dataset | |
| if args.use_validation: | |
| validation_dataset = MegaDepthDataset( | |
| scene_list_path="megadepth_utils/valid_scenes.txt", | |
| scene_info_path=args.scene_info_path, | |
| base_path=args.dataset_path, | |
| train=False, | |
| preprocessing=args.preprocessing, | |
| pairs_per_scene=25, | |
| ) | |
| validation_dataloader = DataLoader( | |
| validation_dataset, batch_size=args.batch_size, num_workers=args.num_workers | |
| ) | |
| training_dataset = MegaDepthDataset( | |
| scene_list_path="megadepth_utils/train_scenes.txt", | |
| scene_info_path=args.scene_info_path, | |
| base_path=args.dataset_path, | |
| preprocessing=args.preprocessing, | |
| ) | |
| training_dataloader = DataLoader( | |
| training_dataset, batch_size=args.batch_size, num_workers=args.num_workers | |
| ) | |
| # Define epoch function | |
| def process_epoch( | |
| epoch_idx, | |
| model, | |
| loss_function, | |
| optimizer, | |
| dataloader, | |
| device, | |
| log_file, | |
| args, | |
| train=True, | |
| ): | |
| epoch_losses = [] | |
| torch.set_grad_enabled(train) | |
| progress_bar = tqdm(enumerate(dataloader), total=len(dataloader)) | |
| for batch_idx, batch in progress_bar: | |
| if train: | |
| optimizer.zero_grad() | |
| batch["train"] = train | |
| batch["epoch_idx"] = epoch_idx | |
| batch["batch_idx"] = batch_idx | |
| batch["batch_size"] = args.batch_size | |
| batch["preprocessing"] = args.preprocessing | |
| batch["log_interval"] = args.log_interval | |
| try: | |
| loss = loss_function(model, batch, device, plot=args.plot) | |
| except NoGradientError: | |
| continue | |
| current_loss = loss.data.cpu().numpy()[0] | |
| epoch_losses.append(current_loss) | |
| progress_bar.set_postfix(loss=("%.4f" % np.mean(epoch_losses))) | |
| if batch_idx % args.log_interval == 0: | |
| log_file.write( | |
| "[%s] epoch %d - batch %d / %d - avg_loss: %f\n" | |
| % ( | |
| "train" if train else "valid", | |
| epoch_idx, | |
| batch_idx, | |
| len(dataloader), | |
| np.mean(epoch_losses), | |
| ) | |
| ) | |
| if train: | |
| loss.backward() | |
| optimizer.step() | |
| log_file.write( | |
| "[%s] epoch %d - avg_loss: %f\n" | |
| % ("train" if train else "valid", epoch_idx, np.mean(epoch_losses)) | |
| ) | |
| log_file.flush() | |
| return np.mean(epoch_losses) | |
| # Create the checkpoint directory | |
| if os.path.isdir(args.checkpoint_directory): | |
| print("[Warning] Checkpoint directory already exists.") | |
| else: | |
| os.mkdir(args.checkpoint_directory) | |
| # Open the log file for writing | |
| if os.path.exists(args.log_file): | |
| print("[Warning] Log file already exists.") | |
| log_file = open(args.log_file, "a+") | |
| # Initialize the history | |
| train_loss_history = [] | |
| validation_loss_history = [] | |
| if args.use_validation: | |
| validation_dataset.build_dataset() | |
| min_validation_loss = process_epoch( | |
| 0, | |
| model, | |
| loss_function, | |
| optimizer, | |
| validation_dataloader, | |
| device, | |
| log_file, | |
| args, | |
| train=False, | |
| ) | |
| # Start the training | |
| for epoch_idx in range(1, args.num_epochs + 1): | |
| # Process epoch | |
| training_dataset.build_dataset() | |
| train_loss_history.append( | |
| process_epoch( | |
| epoch_idx, | |
| model, | |
| loss_function, | |
| optimizer, | |
| training_dataloader, | |
| device, | |
| log_file, | |
| args, | |
| ) | |
| ) | |
| if args.use_validation: | |
| validation_loss_history.append( | |
| process_epoch( | |
| epoch_idx, | |
| model, | |
| loss_function, | |
| optimizer, | |
| validation_dataloader, | |
| device, | |
| log_file, | |
| args, | |
| train=False, | |
| ) | |
| ) | |
| # Save the current checkpoint | |
| checkpoint_path = os.path.join( | |
| args.checkpoint_directory, "%s.%02d.pth" % (args.checkpoint_prefix, epoch_idx) | |
| ) | |
| checkpoint = { | |
| "args": args, | |
| "epoch_idx": epoch_idx, | |
| "model": model.state_dict(), | |
| "optimizer": optimizer.state_dict(), | |
| "train_loss_history": train_loss_history, | |
| "validation_loss_history": validation_loss_history, | |
| } | |
| torch.save(checkpoint, checkpoint_path) | |
| if args.use_validation and validation_loss_history[-1] < min_validation_loss: | |
| min_validation_loss = validation_loss_history[-1] | |
| best_checkpoint_path = os.path.join( | |
| args.checkpoint_directory, "%s.best.pth" % args.checkpoint_prefix | |
| ) | |
| shutil.copy(checkpoint_path, best_checkpoint_path) | |
| # Close the log file | |
| log_file.close() | |