train.py

from tqdm.autonotebook import tqdm
import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchmetrics.classification import MulticlassAccuracy, MulticlassF1Score
from models.classifiers.nn_classifiers.nn_classifier import NNClassifier
from models.loss_functions import GeneralCrossEntropy
from utils.data_utils.tsptw_dataset import TSPTWDataloader
from utils.data_utils.pctsp_dataset import PCTSPDataloader
from utils.data_utils.pctsptw_dataset import PCTSPTWDataloader
from utils.data_utils.cvrp_dataset import CVRPDataloader
from utils.utils import set_device, count_trainable_params, batched_bincount, fix_seed

def main(args):
    #---------------
    # seed settings
    #---------------
    fix_seed(args.seed)

    #--------------
    # gpu settings
    #--------------
    use_cuda, device = set_device(args.gpu)

    #-------------------
    # model & optimizer
    #-------------------
    num_classes = 3 if args.problem == "pctsptw" else 2
    model = NNClassifier(problem=args.problem,
                         node_enc_type=args.node_enc_type,
                         edge_enc_type=args.edge_enc_type,
                         dec_type=args.dec_type,
                         emb_dim=args.emb_dim,
                         num_enc_mlp_layers=args.num_enc_mlp_layers,
                         num_dec_mlp_layers=args.num_dec_mlp_layers,
                         num_classes=num_classes,
                         dropout=args.dropout,
                         pos_encoder=args.pos_encoder)
    is_sequential = model.is_sequential
    if use_cuda:
        model.to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)

    # count number of trainable parameters
    num_trainable_params = count_trainable_params(model)
    print(f"num_trainable_params: {num_trainable_params}")
    with open(f"{args.model_checkpoint_path}/num_trainable_params.dat", "w") as f:
        f.write(str(num_trainable_params))

    # loss function
    if not is_sequential:
        assert args.loss_function != "seq_cbce", "Non-sequential model does not support the loss funtion: seq_cbce"
    loss_func = GeneralCrossEntropy(weight_type=args.loss_function, beta=args.cb_beta, is_sequential=is_sequential)

    #---------
    # dataset
    #---------
    if args.problem == "tsptw":
        train_dataset = TSPTWDataloader(args.train_dataset_path, sequential=is_sequential, parallel=args.parallel, num_cpus=args.num_cpus)
        if args.valid_dataset_path is not None:
            valid_dataset = TSPTWDataloader(args.valid_dataset_path, sequential=is_sequential, parallel=args.parallel, num_cpus=args.num_cpus)
    elif args.problem == "pctsp":
        train_dataset = PCTSPDataloader(args.train_dataset_path, sequential=is_sequential, parallel=args.parallel, num_cpus=args.num_cpus)
        if args.valid_dataset_path is not None:
            valid_dataset = PCTSPDataloader(args.valid_dataset_path, sequential=is_sequential, parallel=args.parallel, num_cpus=args.num_cpus)
    elif args.problem == "pctsptw":
        train_dataset = PCTSPTWDataloader(args.train_dataset_path, sequential=is_sequential, parallel=args.parallel, num_cpus=args.num_cpus)
        if args.valid_dataset_path is not None:
            valid_dataset = PCTSPTWDataloader(args.valid_dataset_path, sequential=is_sequential, parallel=args.parallel, num_cpus=args.num_cpus)
    elif args.problem == "cvrp":
        train_dataset = CVRPDataloader(args.train_dataset_path, sequential=is_sequential, parallel=args.parallel, num_cpus=args.num_cpus)
        if args.valid_dataset_path is not None:
            valid_dataset = CVRPDataloader(args.valid_dataset_path, sequential=is_sequential, parallel=args.parallel, num_cpus=args.num_cpus)
    else:
        raise NotImplementedError

    #------------
    # dataloader
    #------------
    if is_sequential:
        def pad_seq_length(batch):
            data = {}
            for key in batch[0].keys():
                padding_value = True if key == "mask" else 0.0
                # post-padding
                data[key] = torch.nn.utils.rnn.pad_sequence([d[key] for d in batch], batch_first=True, padding_value=padding_value)
            pad_mask = torch.nn.utils.rnn.pad_sequence([torch.full((d["mask"].size(0), ), True) for d in batch], batch_first=True, padding_value=False)
            data.update({"pad_mask": pad_mask})
            return data
        collate_fn = pad_seq_length
    else:
        collate_fn = None
    train_dataloader = DataLoader(train_dataset,
                                  batch_size=args.batch_size,
                                  shuffle=True,
                                  collate_fn=collate_fn,
                                  num_workers=args.num_workers)
    if args.valid_dataset_path is not None:
        valid_dataloader = DataLoader(valid_dataset,
                                      batch_size=args.batch_size,
                                      shuffle=False,
                                      collate_fn=collate_fn,
                                      num_workers=args.num_workers)

    #---------
    # metrics
    #---------
    macro_accuracy = MulticlassF1Score(num_classes=num_classes, average="macro")
    if use_cuda:
        macro_accuracy.to(device)

    #---------------
    # training loop
    #---------------
    best_valid_accuracy = 0.0
    model.train()
    with tqdm(range(args.epochs + 1)) as tq1:
        for epoch in tq1:
            #--------------------------
            # save the current weights
            #--------------------------
            # print(f"Epoch {epoch}: saving a model to {args.model_checkpoint_path}/model_epoch{epoch}.pth...", end="", flush=True)
            torch.save(model.cpu().state_dict(), f"{args.model_checkpoint_path}/model_epoch{epoch}.pth")
            model.to(device)
            # print("done.")

            #------------
            # validation
            #------------
            model.eval()
            with torch.no_grad():
                tq1.set_description(f"Epoch {epoch}")
                # check train accuracy
                for data in train_dataloader:
                    if use_cuda:
                        data = {key: value.to(device) for key, value in data.items()}
                    probs = model(data)
                    if is_sequential:
                        mask = data["pad_mask"].view(-1) # [batch_size x max_seq_length] -> [(batch_size*max_seq_length)]
                        macro_accuracy(probs.argmax(-1).view(-1)[mask], data["labels"].view(-1)[mask])
                    else:
                        macro_accuracy(probs.argmax(-1).view(-1), data["labels"].view(-1))
                train_macro_accuracy = macro_accuracy.compute()
                # print(f"Epoch {epoch}: Train_accuracy={total_macro_accuracy}", flush=True)
                macro_accuracy.reset()

                # check valid accuracy
                if args.valid_dataset_path is not None:
                    for data in valid_dataloader:
                        if use_cuda:
                            data = {key: value.to(device) for key, value in data.items()}
                        probs = model(data)
                        if is_sequential:
                            mask = data["pad_mask"].view(-1) # [batch_size x max_seq_length] -> [(batch_size*max_seq_length)]
                            macro_accuracy(probs.argmax(-1).view(-1)[mask], data["labels"].view(-1)[mask])
                        else:
                            macro_accuracy(probs.argmax(-1).view(-1), data["labels"].view(-1))
                    valid_macro_accuracy = macro_accuracy.compute()
                    # print(f"Epoch {epoch}: Valid_accuracy={total_macro_accuracy}", flush=True)
                    macro_accuracy.reset()
            model.train()
            tq1.set_postfix(Train_accuracy=train_macro_accuracy.item(), Valid_accuracy=valid_macro_accuracy.item())
            
            # update the best epoch
            if valid_macro_accuracy >= best_valid_accuracy:
                best_valid_accuracy = valid_macro_accuracy
                with open(f"{args.model_checkpoint_path}/best_epoch.dat", "w") as f:
                    f.write(str(epoch))

            #--------------------
            # update the weights
            #--------------------
            if epoch < args.epochs:
                with tqdm(train_dataloader, leave=False) as tq:
                    tq.set_description(f"Epoch {epoch}")
                    for data in tq:
                        if use_cuda:
                            data = {key: value.to(device) for key, value in data.items()}
                        out = model(data)
                        if is_sequential:
                            loss = loss_func(out, data["labels"], data["pad_mask"])
                        else:
                            loss = loss_func(out, data["labels"])
                        # if is_sequential:
                        #     # mask = data["pad_mask"].view(-1) # [batch_size x max_seq_length] -> [(batch_size*max_seq_length)]
                        #     # # out = out.view(-1, out.size(-1))
                        #     # # bincount = data["labels"].view(-1)[mask].bincount()
                        #     # # weight = bincount.min() / bincount
                        #     # # loss = F.nll_loss(out[mask], data["labels"].view(-1)[mask], weight=weight)
                        #     # bin = batched_bincount(data["labels"].T, 1, out.size(-1)) # [max_seq_length x num_classes]
                        #     # bin_max, _ = bin.max(-1)
                        #     # weight = bin_max[:, None] / (bin + 1e-8) 
                        #     # weight = weight / weight.max(-1, keepdim=True)[0]
                        #     # # weight = (1 - beta) / (1 - beta**bin)
                        #     # # print(weight)
                        #     # loss = 0.0 # torch.FloatTensor([0.0]).to(device)
                        #     # for seq_no in range(weight.size(0)):
                        #     #     loss += F.nll_loss(out[:, seq_no], data["labels"][:, seq_no], weight=weight[seq_no])
                        # else:
                        #     bincount = data["labels"].view(-1).bincount()
                        #     weight = (1 - beta) / (1 - beta**bincount) 
                        #     loss = F.nll_loss(out, data["labels"].squeeze(-1), weight=weight)
                        optimizer.zero_grad()
                        loss.backward()
                        optimizer.step()
                        tq.set_postfix(Loss=loss.item())

if __name__ == "__main__":
    import datetime
    import json
    import os
    import argparse
    now = datetime.datetime.now()
    parser = argparse.ArgumentParser()
    # general settings
    parser.add_argument("-p", "--problem", default="tsptw", type=str, help="Problem type: [tsptw, cvrptw]")
    parser.add_argument("--gpu", default=-1, type=int, help="Used GPU Number: gpu=-1 indicates using cpu")
    parser.add_argument("--num_workers", default=4, type=int, help="Number of workers in dataloader")
    parser.add_argument("-s", "--seed", type=int, default=1234, help="Random seed for reproductivity")
    # data setting
    parser.add_argument("-train", "--train_dataset_path", type=str, help="Path to a read file", required=True)
    parser.add_argument("-valid", "--valid_dataset_path", type=str, default=None)
    parser.add_argument("--parallel", action="store_true")
    parser.add_argument("--num_cpus", type=int, default=4)
    # training settings
    parser.add_argument("-e", "--epochs", default=100, type=int, help="Number of epochs")
    parser.add_argument("-b", "--batch_size", default=256, type=int, help="Batch size")
    parser.add_argument("--lr", default=0.001, type=float, help="Learning rate")
    parser.add_argument("--cb_beta", default=0.99)
    # parser.add_argument("--valid_interval", default=1, type=int, help="interval outputting intermidiate test accuracy")
    # parser.add_argument("--model_save_interval", type=int, default=1)
    parser.add_argument("--model_checkpoint_path", type=str, default=f"checkpoints/model_{now.strftime('%Y%m%d_%H%M%S')}")
    # model settings
    parser.add_argument("-loss", "--loss_function", type=str, default="seq_cbce", help="[seq_cbce, cbce, wce, ce]")
    parser.add_argument("-node_enc", "--node_enc_type", type=str, default="mlp")
    parser.add_argument("-edge_enc", "--edge_enc_type", type=str, default="attn")
    parser.add_argument("-dec", "--dec_type", type=str, default="lstm")
    parser.add_argument("-pe", "--pos_encoder", type=str, default="sincos")
    parser.add_argument("--emb_dim", type=int, default=128)
    parser.add_argument("--num_enc_mlp_layers", type=int, default=2)
    parser.add_argument("--num_dec_mlp_layers", type=int, default=3)
    parser.add_argument("--dropout", default=0.0, type=float, help="Dropout probability")
    args = parser.parse_args()

    os.makedirs(args.model_checkpoint_path, exist_ok=True)
    with open(f'{args.model_checkpoint_path}/cmd_args.dat', 'w') as f:
        json.dump(args.__dict__, f, indent=2)

    main(args)