Shape-as-Point/train.py

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