train.py (10870B)
1 from dataset import TextGraphDataset 2 from device import device 3 from model import TextRGCN 4 from torch_geometric.loader import DataLoader 5 from tqdm import tqdm 6 import argparse 7 import jsonlines 8 import logging 9 import math 10 import numpy as np 11 import random 12 import torch 13 import torch.nn as nn 14 from utils import set_seed 15 import wandb 16 17 18 def parse_args(): 19 parser = argparse.ArgumentParser(description="Compositional reasoning over KGs and text using LM+GNNs") 20 parser.add_argument("--batch_size", type=int, default=16, help="The batch size to use during training.") 21 parser.add_argument("--weight_decay", type=float, default=0.001, help="The weight decay to use during training.") 22 parser.add_argument("--model_name", type=str, default="bert-base-uncased", help="The pretrained model to use") 23 parser.add_argument("--scheduler", type=str, default="constant", help="The lr scheduler protocol to use") 24 parser.add_argument("--train_path", type=str, default=None, help="The question dataset file for train") 25 parser.add_argument("--val_path", type=str, default=None, help="The question dataset file for val") 26 parser.add_argument("--output_path", type=str, default=None, help="The output predictions path") 27 parser.add_argument("--epochs", type=int, default=3, help="The number of epochs.") 28 parser.add_argument("--gnn_layers", type=int, default=5, help="The number of layers in the GNN model.") 29 parser.add_argument("--gnn_hidden", type=int, default=1024, help="The hidden size of the GNN model.") 30 parser.add_argument("--static", action='store_true', help="Use static entity embeddings") 31 parser.add_argument("--unidirectional", action='store_true', help="Use deep unidirectional mode") 32 parser.add_argument("--bidirectional", action='store_true', help="Use deep unidirectional mode") 33 parser.add_argument("--debug", action='store_true', help="Trigger debug mode") 34 parser.add_argument("--GNN_only", action='store_true', help="Trigger GNN only ablation") 35 parser.add_argument("--LM_only", action='store_true', help="Trigger LM only ablation") 36 parser.add_argument("--lr", type=float, default=3e-5, help="The learning rate).") 37 parser.add_argument("--gnn_lr", type=float, default=3e-5, help="The learning rate).") 38 parser.add_argument("--llm_lr", type=float, default=1e-6, help="The learning rate).") 39 parser.add_argument("--head_lr", type=float, default=1e-5, help="The learning rate).") 40 #parser.add_argument("--seed", type=int, default=42, help="The rng seed") 41 parser.add_argument('--seeds', nargs='+', help='Set of seeds to run with', required=True) 42 parser.add_argument("--device_max_steps", default=200000, type=int, help="Total number of training steps to perform per device") 43 parser.add_argument("--warmup_proportion", default=0.016, type=float, help="Proportion of training to perform linear learning rate warmup for. E.g., 0.1 = 10%% of training.") 44 parser.add_argument("--dropout", type=float, default=0.2, help="The dropout value") 45 args = parser.parse_args() 46 return args 47 48 49 def train(model, loader): 50 total_loss = 0.0 51 correct = 0.0 52 n = 0.0 53 for uids, input_ids, attention_mask, graph_data, head_tail_indicies, special_node_index, special_token_index, label in tqdm(loader, leave=False): 54 optimizer.zero_grad() 55 input_ids = input_ids.to(device) 56 attention_mask = attention_mask.to(device) 57 label = label.to(device) 58 out = model(input_ids, attention_mask, graph_data, head_tail_indicies, special_node_index, special_token_index) 59 loss = criterion(out.squeeze(-1), label) 60 total_loss += loss.item() 61 pred = (nn.functional.sigmoid(out) > 0.5).float().squeeze() 62 correct += (pred == label).sum() 63 n += len(label) 64 loss.backward() 65 optimizer.step() 66 scheduler.step() 67 acc = correct / n 68 loss = total_loss / len(loader) 69 return acc, loss 70 71 72 def val(model, loader): 73 val_loss = 0.0 74 val_preds = [] 75 val_raw_preds = [] 76 val_uids = [] 77 correct = 0.0 78 n = 0.0 79 with torch.no_grad(): 80 for uids, input_ids, attention_mask, graph_data, head_tail_indicies, special_node_index, special_token_index, label in tqdm(val_loader, leave=False): 81 input_ids = input_ids.to(device) 82 attention_mask = attention_mask.to(device) 83 label = label.to(device) 84 out = model(input_ids, attention_mask, graph_data, head_tail_indicies, special_node_index, special_token_index) 85 loss = criterion(out.squeeze(-1), label) 86 val_loss += loss.item() 87 pred = (nn.functional.sigmoid(out) > 0.5).float().squeeze() 88 correct += (pred == label).sum() 89 n += len(label) 90 val_uids += uids.tolist() 91 val_raw_preds += out.cpu().tolist() 92 val_preds += pred.tolist() 93 val_acc = correct / n 94 val_loss = val_loss / len(val_loader) 95 return val_acc, val_loss, val_uids, val_raw_preds, val_preds 96 97 98 if __name__ == '__main__': 99 args = parse_args() 100 logging.basicConfig( 101 format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", 102 datefmt="%m/%d/%Y %H:%M:%S", 103 level=logging.INFO, 104 ) 105 set_seed(0) 106 with jsonlines.open(args.train_path, 'r') as f: 107 train_questions = [o for o in f] 108 with jsonlines.open(args.val_path, 'r') as f: 109 val_questions = [o for o in f] 110 111 train_dataset = TextGraphDataset(train_questions, args=args) 112 val_dataset = TextGraphDataset(val_questions, 113 args=args, 114 head_emb=train_dataset.head_emb, 115 tail_emb=train_dataset.tail_emb, 116 tokenizer=train_dataset.tokenizer) 117 118 logging.info(f'{train_dataset.decode(0)}') 119 logging.info('Starting training...') 120 for seed in args.seeds: 121 logging.info(f'Setting seed {seed}...') 122 set_seed(int(seed)) 123 train_loader = DataLoader(train_dataset, 124 batch_size=args.batch_size, 125 shuffle=True, 126 drop_last=True) 127 val_loader = DataLoader(val_dataset, 128 batch_size=args.batch_size, 129 shuffle=False, 130 drop_last=False) 131 config = { 132 'lr': args.lr, 133 'gnn_lr': args.gnn_lr, 134 'lm_lr': args.llm_lr, 135 'head_lr': args.head_lr, 136 'epochs': args.epochs, 137 'batch_size': args.batch_size, 138 'weight_decay': args.weight_decay, 139 'dropout': args.dropout, 140 'gnn_layers': args.gnn_layers, 141 'scheduler': args.scheduler, 142 'LM_ONLY': args.LM_only, 143 'GNN_ONLY': args.GNN_only, 144 'unidirectional': args.unidirectional, 145 'bidirectional': args.bidirectional, 146 'disjoint': args.static, 147 } 148 wandb.init(project="results", config=config, entity="sondrewo", name=f"{args.output_path}_seed_{seed}") 149 criterion = torch.nn.BCEWithLogitsLoss() 150 model = TextRGCN(args).to(device) 151 if args.bidirectional: 152 model.text_encoder.resize_token_embeddings(len(train_dataset.tokenizer)) 153 154 optimizer = torch.optim.AdamW( 155 [ 156 {"params": model.gcn.parameters(), "lr": args.gnn_lr}, 157 {"params": model.text_encoder.parameters(), "lr": args.llm_lr}, 158 {"params": model.output.parameters(), "lr": args.head_lr}, 159 {"params": model.output2.parameters(), "lr": args.head_lr}, 160 ], 161 lr=args.lr, 162 betas=(0.9, 0.98), 163 eps=1e-6, 164 ) 165 166 args.device_max_steps = args.epochs * len(train_loader) 167 if args.scheduler == "cosine": 168 def cosine_schedule_with_warmup(optimizer, num_warmup_steps: int, num_training_steps: int, min_factor: float): 169 def lr_lambda(current_step): 170 if current_step < num_warmup_steps: 171 return float(current_step) / float(max(1, num_warmup_steps)) 172 progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps)) 173 lr = max(min_factor, min_factor + (1 - min_factor) * 0.5 * (1.0 + math.cos(math.pi * progress))) 174 return lr 175 return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda) 176 scheduler = cosine_schedule_with_warmup(optimizer, int(args.device_max_steps * args.warmup_proportion), args.device_max_steps, 0.1) 177 178 elif args.scheduler == "linear": 179 scheduler = torch.optim.lr_scheduler.ChainedScheduler([ 180 torch.optim.lr_scheduler.LinearLR(optimizer, 181 start_factor=1e-9, 182 end_factor=1.0, 183 total_iters=int(args.device_max_steps * args.warmup_proportion)), 184 torch.optim.lr_scheduler.LinearLR(optimizer, 185 start_factor=1.0, 186 end_factor=1e-9, 187 total_iters=args.device_max_steps) 188 ]) 189 190 elif args.scheduler == "constant": 191 scheduler = torch.optim.lr_scheduler.ConstantLR(optimizer, factor=1.0, total_iters=1) 192 193 for epoch in range(args.epochs): 194 model.train() 195 train_acc, train_loss = train(model, train_loader) 196 wandb.log({'training_accuracy': train_acc, 'training_loss': train_loss}) 197 logging.info(f'Train accuracy: {train_acc}, train.loss: {train_loss}') 198 199 model.eval() 200 val_acc, val_loss, val_uids, val_raw_preds, val_preds = val(model, val_loader) 201 wandb.log({'val_accuracy': val_acc, 'val_loss': val_loss}) 202 logging.info(f'val accuracy: {val_acc}, val.loss: {val_loss}') 203 204 logging.info('Writing final predictions...') 205 final_preds = [] 206 for uid, raw, pred in zip(val_uids, val_raw_preds, val_preds): 207 obj = { 208 'uid': uid, 209 'question': val_dataset.questions[uid], 210 'relations': val_dataset.question_paths[uid], 211 'head_tail': val_dataset.paths[uid], 212 'label': val_dataset.labels[uid], 213 'raw_pred': raw, 214 'pred': pred 215 } 216 final_preds.append(obj) 217 218 with jsonlines.open(f'./results/{args.output_path}_seed_{seed}.jsonl', 'w') as f: 219 f.write_all(final_preds) 220 print('Saved predictions...') 221 wandb.finish()