File size: 11,269 Bytes
d758c99 |
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 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 |
import argparse
import collections
import datetime
import json
import os
import _jsonnet
import attr
import torch
# noinspection PyUnresolvedReferences
from seq2struct import ast_util
# noinspection PyUnresolvedReferences
from seq2struct import datasets
# noinspection PyUnresolvedReferences
from seq2struct import models
# noinspection PyUnresolvedReferences
from seq2struct import optimizers
from seq2struct.utils import registry
from seq2struct.utils import random_state
from seq2struct.utils import saver as saver_mod
# noinspection PyUnresolvedReferences
from seq2struct.utils import vocab
@attr.s
class TrainConfig:
eval_every_n = attr.ib(default=100)
report_every_n = attr.ib(default=100)
save_every_n = attr.ib(default=100)
keep_every_n = attr.ib(default=1000)
batch_size = attr.ib(default=32)
eval_batch_size = attr.ib(default=32)
max_steps = attr.ib(default=100000)
num_eval_items = attr.ib(default=None)
eval_on_train = attr.ib(default=True)
eval_on_val = attr.ib(default=True)
# Seed for RNG used in shuffling the training data.
data_seed = attr.ib(default=None)
# Seed for RNG used in initializing the model.
init_seed = attr.ib(default=None)
# Seed for RNG used in computing the model's training loss.
# Only relevant with internal randomness in the model, e.g. with dropout.
model_seed = attr.ib(default=None)
num_batch_accumulated = attr.ib(default=1)
clip_grad = attr.ib(default=None)
class Logger:
def __init__(self, log_path=None, reopen_to_flush=False):
self.log_file = None
self.reopen_to_flush = reopen_to_flush
if log_path is not None:
os.makedirs(os.path.dirname(log_path), exist_ok=True)
self.log_file = open(log_path, 'a+')
def log(self, msg):
formatted = '[{}] {}'.format(
datetime.datetime.now().replace(microsecond=0).isoformat(),
msg)
print(formatted)
if self.log_file:
self.log_file.write(formatted + '\n')
if self.reopen_to_flush:
log_path = self.log_file.name
self.log_file.close()
self.log_file = open(log_path, 'a+')
else:
self.log_file.flush()
class Trainer:
def __init__(self, logger, config):
if torch.cuda.is_available():
self.device = torch.device('cuda')
else:
self.device = torch.device('cpu')
self.logger = logger
self.train_config = registry.instantiate(TrainConfig, config['train'])
self.data_random = random_state.RandomContext(self.train_config.data_seed)
self.model_random = random_state.RandomContext(self.train_config.model_seed)
self.init_random = random_state.RandomContext(self.train_config.init_seed)
with self.init_random:
# 0. Construct preprocessors
self.model_preproc = registry.instantiate(
registry.lookup('model', config['model']).Preproc,
config['model'],
unused_keys=('name',))
self.model_preproc.load()
# 1. Construct model
self.model = registry.construct('model', config['model'],
unused_keys=('encoder_preproc', 'decoder_preproc'), preproc=self.model_preproc, device=self.device)
self.model.to(self.device)
def train(self, config, modeldir):
# slight difference here vs. unrefactored train: The init_random starts over here. Could be fixed if it was important by saving random state at end of init
with self.init_random:
# We may be able to move optimizer and lr_scheduler to __init__ instead. Empirically it works fine. I think that's because saver.restore
# resets the state by calling optimizer.load_state_dict.
# But, if there is no saved file yet, I think this is not true, so might need to reset the optimizer manually?
# For now, just creating it from scratch each time is safer and appears to be the same speed, but also means you have to pass in the config to train which is kind of ugly.
# TODO: not nice
if config["optimizer"].get("name", None) == 'bertAdamw':
bert_params = list(self.model.encoder.bert_model.parameters())
assert len(bert_params) > 0
non_bert_params = []
for name, _param in self.model.named_parameters():
if "bert" not in name:
non_bert_params.append(_param)
assert len(non_bert_params) + len(bert_params) == len(list(self.model.parameters()))
optimizer = registry.construct('optimizer', config['optimizer'], non_bert_params=non_bert_params, \
bert_params=bert_params)
lr_scheduler = registry.construct( 'lr_scheduler',
config.get('lr_scheduler', {'name': 'noop'}),
param_groups=[optimizer.non_bert_param_group, \
optimizer.bert_param_group])
else:
optimizer = registry.construct('optimizer', config['optimizer'], params=self.model.parameters())
lr_scheduler = registry.construct( 'lr_scheduler',
config.get('lr_scheduler', {'name': 'noop'}),
param_groups=optimizer.param_groups)
# 2. Restore model parameters
saver = saver_mod.Saver(
{"model": self.model, "optimizer": optimizer}, keep_every_n=self.train_config.keep_every_n)
last_step = saver.restore(modeldir, map_location=self.device)
if "pretrain" in config and last_step == 0:
pretrain_config = config["pretrain"]
_path = pretrain_config["pretrained_path"]
_step = pretrain_config["checkpoint_step"]
pretrain_step = saver.restore(_path, step=_step, map_location=self.device, item_keys=["model"])
saver.save(modeldir, pretrain_step) # for evaluating pretrained models
last_step = pretrain_step
# 3. Get training data somewhere
with self.data_random:
train_data = self.model_preproc.dataset('train')
train_data_loader = self._yield_batches_from_epochs(
torch.utils.data.DataLoader(
train_data,
batch_size=self.train_config.batch_size,
shuffle=True,
drop_last=True,
collate_fn=lambda x: x))
train_eval_data_loader = torch.utils.data.DataLoader(
train_data,
batch_size=self.train_config.eval_batch_size,
collate_fn=lambda x: x)
val_data = self.model_preproc.dataset('val')
val_data_loader = torch.utils.data.DataLoader(
val_data,
batch_size=self.train_config.eval_batch_size,
collate_fn=lambda x: x)
# 4. Start training loop
with self.data_random:
for batch in train_data_loader:
# Quit if too long
if last_step >= self.train_config.max_steps:
break
# Evaluate model
if last_step % self.train_config.eval_every_n == 0:
if self.train_config.eval_on_train:
self._eval_model(self.logger, self.model, last_step, train_eval_data_loader, 'train', num_eval_items=self.train_config.num_eval_items)
if self.train_config.eval_on_val:
self._eval_model(self.logger, self.model, last_step, val_data_loader, 'val', num_eval_items=self.train_config.num_eval_items)
# Compute and apply gradient
with self.model_random:
for _i in range(self.train_config.num_batch_accumulated):
if _i > 0: batch = next(train_data_loader)
loss = self.model.compute_loss(batch)
norm_loss = loss / self.train_config.num_batch_accumulated
norm_loss.backward()
if self.train_config.clip_grad:
torch.nn.utils.clip_grad_norm_(optimizer.bert_param_group["params"], \
self.train_config.clip_grad)
optimizer.step()
lr_scheduler.update_lr(last_step)
optimizer.zero_grad()
# Report metrics
if last_step % self.train_config.report_every_n == 0:
self.logger.log('Step {}: loss={:.4f}'.format(last_step, loss.item()))
last_step += 1
# Run saver
if last_step % self.train_config.save_every_n == 0:
saver.save(modeldir, last_step)
# Save final model
saver.save(modeldir, last_step)
@staticmethod
def _yield_batches_from_epochs(loader):
while True:
for batch in loader:
yield batch
@staticmethod
def _eval_model(logger, model, last_step, eval_data_loader, eval_section, num_eval_items=None):
stats = collections.defaultdict(float)
model.eval()
with torch.no_grad():
for eval_batch in eval_data_loader:
batch_res = model.eval_on_batch(eval_batch)
for k, v in batch_res.items():
stats[k] += v
if num_eval_items and stats['total'] > num_eval_items:
break
model.train()
# Divide each stat by 'total'
for k in stats:
if k != 'total':
stats[k] /= stats['total']
if 'total' in stats:
del stats['total']
logger.log("Step {} stats, {}: {}".format(
last_step, eval_section, ", ".join(
"{} = {}".format(k, v) for k, v in stats.items())))
def add_parser():
parser = argparse.ArgumentParser()
parser.add_argument('--logdir', required=True)
parser.add_argument('--config', required=True)
parser.add_argument('--config-args')
args = parser.parse_args()
return args
def main(args):
if args.config_args:
config = json.loads(_jsonnet.evaluate_file(args.config, tla_codes={'args': args.config_args}))
else:
config = json.loads(_jsonnet.evaluate_file(args.config))
if 'model_name' in config:
args.logdir = os.path.join(args.logdir, config['model_name'])
# Initialize the logger
reopen_to_flush = config.get('log', {}).get('reopen_to_flush')
logger = Logger(os.path.join(args.logdir, 'log.txt'), reopen_to_flush)
# Save the config info
with open(os.path.join(args.logdir,
'config-{}.json'.format(
datetime.datetime.now().strftime('%Y%m%dT%H%M%S%Z'))), 'w', encoding='utf8') as f:
json.dump(config, f, sort_keys=True, indent=4, ensure_ascii=False)
logger.log('Logging to {}'.format(args.logdir))
# Construct trainer and do training
trainer = Trainer(logger, config)
trainer.train(config, modeldir=args.logdir)
if __name__ == '__main__':
args = add_parser()
main(args)
|