-
Notifications
You must be signed in to change notification settings - Fork 2
/
ner_gradients.py
executable file
·82 lines (65 loc) · 2.85 KB
/
ner_gradients.py
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
import os
import numpy as np
import torch
import tqdm
from numpy import save
from dataloaders.ner_conll2003 import get_labels, conll2003_get_dataloader
from models.BertSequence import load_bert_sequence_model
from influencer.GD import GD
from torch.autograd import grad
def build_gradient(model, params, dataloader, type_data = None):
for i, data in enumerate(tqdm.tqdm(dataloader)):
# if os.path.exists(os.path.join(dir_checkpoint, f"ck2_{type_data}_gradients", f'grad_{type_data}_{i}')):
# continue
ids = data['input_ids'].to(device, dtype = torch.long)
mask = data['attention_mask'].to(device, dtype = torch.long)
labels = data['labels'].to(device, dtype = torch.long)
loss, tr_logits = model(input_ids=ids, attention_mask=mask, labels=labels, reduction_loss = 'none')
sentence_gradient = [grad(l, params, create_graph=False, retain_graph=True) for l in loss.view(-1)]
lst_gradient_of_sentence = []
for word_gradient in sentence_gradient:
weight, bias = word_gradient
weight = weight.cpu().detach().numpy()
bias = bias.cpu().detach().numpy()
word_gradient_numpy = [weight, bias]
lst_gradient_of_sentence.append(word_gradient_numpy)
# save gradient
torch.save(lst_gradient_of_sentence, os.path.join(dir_checkpoint, f"ck2_{type_data}_gradients", f'grad_{type_data}_{i}'))
if __name__ == '__main__':
# noise_BItags_30sen_30tok
noise_file = os.path.join('data', 'conll2003', 'ner', 'noise_BItags_30sen_30tok.txt')
dir_checkpoint = os.path.join('checkpoints', 'conll2003', 'SEED4_NER_CoNLL2003_noise_BItags_30sen_30tok')
device = 'cuda'
num_labels = len(get_labels())
type_data = "noise"
SEED = 4
name_checkpoint = 'epoch_2.pt' # best.pt
np.random.seed(SEED)
torch.manual_seed(SEED)
torch.cuda.manual_seed(SEED)
if not torch.cuda.is_available() and device == 'cuda':
print('Your device don\'t have cuda, please check or select cpu and retraining')
exit(0)
if device == 'cuda':
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
if torch.cuda.device_count() > 1:
print("Let's use", torch.cuda.device_count(), "GPUs!")
# Load data
noise_dataloader = conll2003_get_dataloader(
file_name=noise_file,
batch_size=1,
mode='test',
num_workers=0
)
# Build model
model = load_bert_sequence_model(os.path.join(dir_checkpoint, name_checkpoint),num_labels=num_labels,device=device)
# Params to get gradients
params = [p for p in model.parameters() if p.requires_grad][-2:]
# get gradients
build_gradient(
model = model,
params = params,
dataloader = noise_dataloader,
type_data = type_data
)