import os abspath = os.path.abspath(__file__) dname = os.path.dirname(abspath) os.chdir(dname) import argparse import shutil import time import numpy as np import torch import torch.optim as optim from torch.utils.tensorboard import SummaryWriter from src import config from src.data import collate_remove_none, worker_init_fn from src.model import Encode2Points from src.training import Trainer from src.utils import AverageMeter, initialize_logger, load_config, load_model_manual np.set_printoptions(precision=4) def main(): parser = argparse.ArgumentParser(description="MNIST toy experiment") parser.add_argument("config", type=str, help="Path to config file.") parser.add_argument("--no_cuda", action="store_true", default=False, help="disables CUDA training") parser.add_argument("--seed", type=int, default=1, metavar="S", help="random seed (default: 1)") args = parser.parse_args() cfg = load_config(args.config, "configs/default.yaml") use_cuda = not args.no_cuda and torch.cuda.is_available() device = torch.device("cuda" if use_cuda else "cpu") cfg["data"]["input_type"] batch_size = cfg["train"]["batch_size"] model_selection_metric = cfg["train"]["model_selection_metric"] # PYTORCH VERSION > 1.0.0 assert float(torch.__version__.split(".")[-3]) > 0 # boiler-plate if cfg["train"]["timestamp"]: cfg["train"]["out_dir"] += "_" + time.strftime("%Y_%m_%d_%H_%M_%S") logger = initialize_logger(cfg) torch.manual_seed(args.seed) np.random.seed(args.seed) shutil.copyfile(args.config, os.path.join(cfg["train"]["out_dir"], "config.yaml")) logger.info("using GPU: " + torch.cuda.get_device_name(0)) # TensorboardX writer tblogdir = os.path.join(cfg["train"]["out_dir"], "tensorboard_log") if not os.path.exists(tblogdir): os.makedirs(tblogdir, exist_ok=True) writer = SummaryWriter(log_dir=tblogdir) inputs = None train_dataset = config.get_dataset("train", cfg) val_dataset = config.get_dataset("val", cfg) vis_dataset = config.get_dataset("vis", cfg) collate_fn = collate_remove_none train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=batch_size, num_workers=cfg["train"]["n_workers"], shuffle=True, collate_fn=collate_fn, worker_init_fn=worker_init_fn, ) val_loader = torch.utils.data.DataLoader( val_dataset, batch_size=1, num_workers=cfg["train"]["n_workers_val"], shuffle=False, collate_fn=collate_remove_none, worker_init_fn=worker_init_fn, ) vis_loader = torch.utils.data.DataLoader( vis_dataset, batch_size=1, num_workers=cfg["train"]["n_workers_val"], shuffle=False, collate_fn=collate_fn, worker_init_fn=worker_init_fn, ) if torch.cuda.device_count() > 1: model = torch.nn.DataParallel(Encode2Points(cfg)).to(device) else: model = Encode2Points(cfg).to(device) n_parameter = sum(p.numel() for p in model.parameters() if p.requires_grad) logger.info("Number of parameters: %d" % n_parameter) # load model try: # load model state_dict = torch.load(os.path.join(cfg["train"]["out_dir"], "model.pt")) load_model_manual(state_dict["state_dict"], model) out = "Load model from iteration %d" % state_dict.get("it", 0) logger.info(out) # load point cloud except: state_dict = dict() metric_val_best = state_dict.get("loss_val_best", np.inf) logger.info(f"Current best validation metric ({model_selection_metric}): {metric_val_best:.8f}") LR = float(cfg["train"]["lr"]) optimizer = optim.Adam(model.parameters(), lr=LR) start_epoch = state_dict.get("epoch", -1) it = state_dict.get("it", -1) trainer = Trainer(cfg, optimizer, device=device) runtime = {} runtime["all"] = AverageMeter() # training loop for epoch in range(start_epoch + 1, cfg["train"]["total_epochs"] + 1): for batch in train_loader: it += 1 start = time.time() loss, loss_each = trainer.train_step(inputs, batch, model) # measure elapsed time end = time.time() runtime["all"].update(end - start) if it % cfg["train"]["print_every"] == 0: log_text = ("[Epoch %02d] it=%d, loss=%.4f") % (epoch, it, loss) writer.add_scalar("train/loss", loss, it) if loss_each is not None: for k, l in loss_each.items(): if l.item() != 0.0: log_text += f" loss_{k}={l.item():.4f}" writer.add_scalar("train/%s" % k, l, it) log_text += (" time={:.3f} / {:.2f}").format(runtime["all"].val, runtime["all"].sum) logger.info(log_text) if (it > 0) & (it % cfg["train"]["visualize_every"] == 0): for i, batch_vis in enumerate(vis_loader): trainer.save(model, batch_vis, it, i) if i >= 4: break logger.info("Saved mesh and pointcloud") # run validation if it > 0 and (it % cfg["train"]["validate_every"]) == 0: eval_dict = trainer.evaluate(val_loader, model) metric_val = eval_dict[model_selection_metric] logger.info(f"Validation metric ({model_selection_metric}): {metric_val:.4f}") for k, v in eval_dict.items(): writer.add_scalar("val/%s" % k, v, it) if -(metric_val - metric_val_best) >= 0: metric_val_best = metric_val logger.info("New best model (loss %.4f)" % metric_val_best) state = {"epoch": epoch, "it": it, "loss_val_best": metric_val_best} state["state_dict"] = model.state_dict() torch.save(state, os.path.join(cfg["train"]["out_dir"], "model_best.pt")) # save checkpoint if (epoch > 0) & (it % cfg["train"]["checkpoint_every"] == 0): state = {"epoch": epoch, "it": it, "loss_val_best": metric_val_best} state["state_dict"] = model.state_dict() torch.save(state, os.path.join(cfg["train"]["out_dir"], "model.pt")) if it % cfg["train"]["backup_every"] == 0: torch.save(state, os.path.join(cfg["train"]["dir_model"], "%04d" % it + ".pt")) logger.info("Backup model at iteration %d" % it) logger.info("Save new model at iteration %d" % it) time.time() if __name__ == "__main__": main()