Spaces:
Sleeping
Sleeping
File size: 6,228 Bytes
63239ac | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 | 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']
|