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()