File size: 3,142 Bytes
2d06dcc |
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 |
import os
import torch
import math
import logging
from transformers import AdamW, get_linear_schedule_with_warmup
from .utils import freeze_bert_parameters, set_allow_growth
from .__init__ import backbones_map
class ModelManager:
def __init__(self, args, data, logger_name = 'Discovery'):
self.logger = logging.getLogger(logger_name)
if args.method in ['KM', 'AG', 'SAE', 'DEC', 'DCN']:
set_allow_growth('0')
else:
self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
def set_optimizer(self, model, num_train_examples, train_batch_size, num_train_epochs, lr, warmup_proportion):
num_train_optimization_steps = int(num_train_examples / train_batch_size) * num_train_epochs
param_optimizer = list(model.named_parameters())
no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']
optimizer_grouped_parameters = [
{'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)], 'weight_decay': 0.01},
{'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}
]
optimizer = AdamW(optimizer_grouped_parameters, lr = lr, correct_bias=False)
num_warmup_steps= int(num_train_examples * num_train_epochs * warmup_proportion / train_batch_size)
scheduler = get_linear_schedule_with_warmup(optimizer,
num_warmup_steps=num_warmup_steps,
num_training_steps=num_train_optimization_steps)
return optimizer ,scheduler
def set_model(self, args, data, pattern, freeze_parameters = True):
backbone = backbones_map[args.backbone]
if pattern == 'bert':
if hasattr(backbone, 'from_pretrained'):
model = backbone.from_pretrained(args.pretrained_bert_model, args = args)
else:
model = backbone(args)
if freeze_parameters:
self.logger.info('Freeze all parameters but the last layer for efficiency')
model = freeze_bert_parameters(model)
model.to(self.device)
return model
elif args.setting == 'unsupervised':
if pattern == 'glove':
self.logger.info("Building GloVe (D=300)...")
gev = backbone(data.dataloader.embedding_matrix, data.dataloader.index_word, data.dataloader.train_data)
emb_train = gev.transform(data.dataloader.train_data, method='mean')
emb_test = gev.transform(data.dataloader.test_data, method='mean')
self.logger.info('Building finished!')
return emb_train, emb_test
elif pattern == 'sae':
self.logger.info("Building TF-IDF Vectors...")
sae = backbone(data.dataloader.tfidf_train.shape[1])
return sae
|