Spaces:
Sleeping
Sleeping
| 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) |