import os import numpy as np import torch import torch.nn as nn from collections import deque from torch.utils.data import DataLoader from tqdm import tqdm from typing import Optional, Mapping from deepxml.evaluation import get_p_1, get_p_3, get_p_5, get_n_1, get_n_3, get_n_5 from deepxml.optimizers import DenseSparseAdam from deepxml.data_utils import truncate_text class Model(object): def __init__(self, network, model_path, mode, graph_hierarchy=None ,reg=False, gradient_clip_value=5.0, device_ids=None, **kwargs): self.model = nn.DataParallel(network(graph_hierarchy=graph_hierarchy, **kwargs), device_ids=device_ids) self.loss_fn = nn.BCEWithLogitsLoss() self.model_path, self.state = model_path, {} os.makedirs(os.path.split(self.model_path)[0], exist_ok=True) self.gradient_clip_value, self.gradient_norm_queue = gradient_clip_value, deque([np.inf], maxlen=5) self.optimizer = None # self.load_model() self.reg = reg if mode == 'train' and reg: self.hierarchy = graph_hierarchy["hierarchy"] self.lambda1 = 1e-8 self.lambda2 = 1e-10 def train_step(self, train_data: torch.Tensor, train_y: torch.Tensor): self.optimizer.zero_grad() self.model.train() scores = self.model(train_data) # scores = scores.view(train_y.shape[0], train_y.shape[1]) loss = self.loss_fn(scores, train_y) if self.reg: # Output Regularization probs = torch.sigmoid(scores) regs = torch.zeros(len(probs), len(self.hierarchy)).cuda() for idx, tup in enumerate(self.hierarchy): p = tup[0] c = tup[1] regs[:,idx] = probs[:,c] - probs[:,p] loss += self.lambda1 * torch.sum(nn.functional.relu(regs)).item() # Parameter Regularization # weights = self.model.module.plaincls.out_mesh_dstrbtn.weight # regs = torch.zeros(len(weights[0]), len(self.hierarchy)).cuda() # for idx, tup in enumerate(self.hierarchy): # p = tup[0] # c = tup[1] # regs[:,idx] = weights[p] - weights[c] # loss += self.lambda2 * 1/2 * torch.norm(regs, p=2) ** 2 loss.backward() self.clip_gradient() self.optimizer.step(closure=None) return loss.item() def predict_step(self, data_x: torch.Tensor, k: int): self.model.eval() with torch.no_grad(): scores, labels = torch.topk(self.model(data_x), k) return torch.sigmoid(scores).cpu(), labels.cpu() def get_optimizer(self, **kwargs): self.optimizer = DenseSparseAdam(self.model.parameters(), **kwargs) def train(self, train_loader: DataLoader, valid_loader: DataLoader, opt_params: Optional[Mapping] = None, nb_epoch=100, step=100, k=5, early=100, verbose=True, swa_warmup=None, **kwargs): self.get_optimizer(**({} if opt_params is None else opt_params)) global_step, best_n5, e = 0, 0.0, 0 print_loss = 0.0 for epoch_idx in range(nb_epoch): if epoch_idx == swa_warmup: self.swa_init() for i, (train_x, train_y) in enumerate(train_loader, 1): global_step += 1 loss = self.train_step(train_x, train_y.cuda()) print_loss += loss if global_step % step == 0: self.swa_step() self.swap_swa_params() labels = [] valid_loss = 0.0 self.model.eval() with torch.no_grad(): for (valid_x, valid_y) in valid_loader: logits = self.model(valid_x) # logits = logits.view(valid_y.shape[0], valid_y.shape[1]) valid_loss += self.loss_fn(logits, valid_y.cuda()).item() scores, tmp = torch.topk(logits, k) labels.append(tmp.cpu()) valid_loss /= len(valid_loader) labels = np.concatenate(labels) targets = valid_loader.dataset.data_y[:len(labels),:] p1, p3, p5, n3, n5 = get_p_1(labels, targets), get_p_3(labels, targets), get_p_5(labels, targets), get_n_3(labels, targets), get_n_5(labels, targets) if n5 >= best_n5: self.save_model(True) best_n5, e = n5, 0 else: e += 1 if early is not None and e > early: return self.swap_swa_params() if verbose: log_msg = '%d %d train loss: %.7f valid loss: %.7f P@1: %.5f P@3: %.5f P@5: %.5f N@3: %.5f N@5: %.5f early stop: %d' % \ (epoch_idx, i * train_loader.batch_size, print_loss / step, valid_loss, round(p1, 5), round(p3, 5), round(p5, 5), round(n3, 5), round(n5, 5), e) logger.info(log_msg) print_loss = 0.0 fh = open('best.txt', 'a', encoding='utf-8') fh.write(log_msg) fh.write('\n') fh.close() def predict(self, data_x, desc='Predict', **kwargs): self.load_model() self.model.eval() with torch.no_grad(): scores= self.model(data_x) return torch.sigmoid(scores) def save_model(self, last_epoch): if not last_epoch: return for trial in range(5): try: torch.save(self.model.module.state_dict(), self.model_path) break except: print('saving failed') def load_model(self): self.model.module.load_state_dict(torch.load(self.model_path, map_location=torch.device('cpu'))) def clip_gradient(self): if self.gradient_clip_value is not None: max_norm = max(self.gradient_norm_queue) total_norm = torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm * self.gradient_clip_value) self.gradient_norm_queue.append(min(total_norm, max_norm * 2.0, 1.0)) if total_norm > max_norm * self.gradient_clip_value: logger.warn(F'Clipping gradients with total norm {total_norm} ' F'and max norm {max_norm}') def swa_init(self): if 'swa' not in self.state: logger.info('SWA Initializing') swa_state = self.state['swa'] = {'models_num': 1} for n, p in self.model.named_parameters(): if p.requires_grad: swa_state[n] = p.data.cpu().detach() def swa_step(self): if 'swa' in self.state: swa_state = self.state['swa'] swa_state['models_num'] += 1 beta = 1.0 / swa_state['models_num'] with torch.no_grad(): for n, p in self.model.named_parameters(): if p.requires_grad: swa_state[n].mul_(1.0 - beta).add_(beta, p.data.cpu()) def swap_swa_params(self): if 'swa' in self.state: swa_state = self.state['swa'] for n, p in self.model.named_parameters(): if p.requires_grad: p.data, swa_state[n] = swa_state[n].cuda(), p.data.cpu() def disable_swa(self): if 'swa' in self.state: del self.state['swa']