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)
orig_data = registry.construct('dataset', config['data']["train"])
self.model.load_orig_data(orig_data)
def construct_optimizer_and_lr_scheduler(self, config):
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)
return optimizer, lr_scheduler
def __init__(self, preproc, device, encoder, decoder):
super().__init__()
self.preproc = preproc
self.encoder = registry.construct(
'encoder', encoder, device=device, preproc=preproc.enc_preproc)
self.decoder = registry.construct(
'decoder', decoder, device=device, preproc=preproc.dec_preproc)
if getattr(self.encoder, 'batched'):
self.compute_loss = self._compute_loss_enc_batched
else:
self.compute_loss = self._compute_loss_unbatched
def __init__(self,
save_path,
min_freq=3,
max_count=5000,
include_table_name_in_column=True,
word_emb=None,
count_tokens_in_word_emb_for_vocab=False,
fix_issue_16_primary_keys=False,
compute_sc_link=False,
compute_cv_link=False,
db_path=None):
if word_emb is None:
self.word_emb = None
else:
self.word_emb = registry.construct('word_emb', word_emb)
self.data_dir = os.path.join(save_path, 'enc')
self.include_table_name_in_column = include_table_name_in_column
self.count_tokens_in_word_emb_for_vocab = count_tokens_in_word_emb_for_vocab
self.fix_issue_16_primary_keys = fix_issue_16_primary_keys
self.compute_sc_link = compute_sc_link
self.compute_cv_link = compute_cv_link
self.texts = collections.defaultdict(list)
self.db_path = db_path
self.vocab_builder = vocab.VocabBuilder(min_freq, max_count)
self.vocab_path = os.path.join(save_path, 'enc_vocab.json')
self.vocab_word_freq_path = os.path.join(save_path,
'enc_word_freq.json')
self.vocab = None
self.counted_db_ids = set()
self.preprocessed_schemas = {}
def infer(self, model, output_path, args):
output = open(output_path, 'w')
with torch.no_grad():
if args.mode == 'infer':
orig_data = registry.construct(
'dataset', self.config['data'][args.section])
preproc_data = self.model_preproc.dataset(args.section)
if args.limit:
sliced_orig_data = itertools.islice(orig_data, args.limit)
sliced_preproc_data = itertools.islice(
preproc_data, args.limit)
else:
sliced_orig_data = orig_data
sliced_preproc_data = preproc_data
assert len(orig_data) == len(preproc_data)
self._inner_infer(model, args.beam_size, args.output_history,
sliced_orig_data, sliced_preproc_data,
output, args.use_heuristic)
elif args.mode == 'debug':
data = self.model_preproc.dataset(args.section)
if args.limit:
sliced_data = itertools.islice(data, args.limit)
else:
sliced_data = data
self._debug(model, sliced_data, output)
def __init__(self, logger, config):
self.config = config
if torch.cuda.is_available():
self.device = torch.device('cuda')
else:
self.device = torch.device('cpu')
self.logger = logger
self.finetune_config = registry.instantiate(FineTuneConfig,
config['finetune'])
self.model_random = random_state.RandomContext(
self.finetune_config.model_seed)
self.init_random = random_state.RandomContext(
self.finetune_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 __init__(self,
grammar,
save_path,
min_freq=3,
max_count=5000,
use_seq_elem_rules=False):
self.grammar = registry.construct('grammar', grammar)
self.ast_wrapper = self.grammar.ast_wrapper
self.vocab_path = os.path.join(save_path, 'dec_vocab.json')
self.observed_productions_path = os.path.join(
save_path, 'observed_productions.json')
self.grammar_rules_path = os.path.join(save_path, 'grammar_rules.json')
self.data_dir = os.path.join(save_path, 'dec')
self.vocab_builder = vocab.VocabBuilder(min_freq, max_count)
self.use_seq_elem_rules = use_seq_elem_rules
self.items = collections.defaultdict(list)
self.sum_type_constructors = collections.defaultdict(set)
self.field_presence_infos = collections.defaultdict(set)
self.seq_lengths = collections.defaultdict(set)
self.primitive_types = set()
self.vocab = None
self.all_rules = None
self.rules_mask = None
def preprocess(self):
self.model_preproc.clear_items()
for section in self.config['data']:
data = registry.construct('dataset', self.config['data'][section])
for item in tqdm.tqdm(data, desc=f"{section} section", dynamic_ncols=True):
to_add, validation_info = self.model_preproc.validate_item(item, section)
if to_add:
self.model_preproc.add_item(item, section, validation_info)
self.model_preproc.save()
def __init__(self, preproc, device, encoder, decoder):
super().__init__()
self.preproc = preproc
self._device = device
self.encoder = registry.construct(
'encoder', encoder, device=device, preproc=preproc.enc_preproc)
self.decoder = registry.construct(
'decoder', decoder, device=device, preproc=preproc.dec_preproc)
if self.encoder.use_discourse_level_lstm:
self.discourse_lstms = create_multilayer_lstm_params(1, self.encoder.enc_hidden_size, self.encoder.enc_hidden_size / 2)
self.initial_discourse_state = add_params(tuple([self.encoder.enc_hidden_size / 2]))
if self.encoder.use_utterance_attention:
self.utterance_attention_module = Attention(self.encoder.enc_hidden_size, self.encoder.enc_hidden_size, self.encoder.enc_hidden_size)
if getattr(self.encoder, 'batched'):
self.compute_loss = self._compute_loss_enc_batched ##encode有batched==True, sparc时为false
else:
self.compute_loss = self._compute_loss_unbatched #走这个
def load_model(self, logdir, step):
'''Load a model (identified by the config used for construction) and return it'''
# 1. Construct model 创建model
model = registry.construct('model', self.config['model'], preproc=self.model_preproc, device=self.device)
model.to(self.device)
model.eval()
# 2. Restore its parameters
saver = saver_mod.Saver({"model": model})
last_step = saver.restore(logdir, step=step, map_location=self.device, item_keys=["model"])
if not last_step:
raise Exception(f"Attempting to infer on untrained model in {logdir}, step={step}")
return model
def preprocess(self):
self.model_preproc.clear_items()
for section in self.config['data']:
data = registry.construct('dataset', self.config['data'][section])
# <ratsql.datasets.spider.SpiderDataset object at 0x7f86d5dc8690>
# DB connections is done with the construction of the dataset
for item in tqdm.tqdm(data,
desc=f"{section} section",
dynamic_ncols=True):
to_add, validation_info = self.model_preproc.validate_item(
item, section)
if to_add:
self.model_preproc.add_item(item, section, validation_info)
self.model_preproc.save()
def compute_metrics(config_path,
config_args,
section,
inferred_path,
logdir=None,
infer_type='inferred_code'):
if config_args:
config = json.loads(
_jsonnet.evaluate_file(config_path,
tla_codes={'args': config_args}))
else:
config = json.loads(_jsonnet.evaluate_file(config_path))
if 'model_name' in config and logdir:
logdir = os.path.join(logdir, config['model_name'])
if logdir:
inferred_path = inferred_path.replace('__LOGDIR__', logdir)
inferred = open(inferred_path)
data = registry.construct('dataset', config['data'][section])
metrics = data.Metrics(data)
inferred_lines = list(inferred)
if len(inferred_lines) < len(data):
raise Exception(
f'Not enough inferred: {len(inferred_lines)} vs {len(data)}')
for line in tqdm.tqdm(inferred_lines):
infer_results = json.loads(line)
if infer_results['beams']:
inferred_code = infer_results['beams'][0][infer_type]
else:
inferred_code = None
if 'index' in infer_results:
metrics.add(data[infer_results['index']], inferred_code)
else:
metrics.add(None,
inferred_code,
obsolete_gold_code=infer_results['gold_code'])
return logdir, metrics.finalize()
def compute_metrics(config_path,
config_args,
section,
inferred_path,
logdir=None):
if config_args:
config = json.loads(
_jsonnet.evaluate_file(config_path,
tla_codes={'args': config_args}))
else:
config = json.loads(_jsonnet.evaluate_file(config_path))
if 'model_name' in config and logdir:
logdir = os.path.join(logdir, config['model_name'])
if logdir:
inferred_path = inferred_path.replace('__LOGDIR__', logdir)
inferred = open(inferred_path)
data = registry.construct('dataset', config['data'][section])
metrics = data.Metrics(data)
data_len = 0
for interaction in data:
data_len += len(interaction.utterances)
inferred_lines = list(inferred) #预测的interaction数
if len(inferred_lines) < data_len: #如果小于data的个数
raise Exception(
f'Not enough inferred: {len(inferred_lines)} vs {len(data)}')
for line in inferred_lines:
infer_results = json.loads(line)
inferred_code = infer_results["beams"][0]["inferred_code"]
utterance_index = infer_results["beams"][0]["utterance_index"]
schema_id = data[infer_results["interaction_index"]].schema.db_id
ori_query = data[
infer_results["interaction_index"]].querys[utterance_index]
metrics.add(schema_id, ori_query, inferred_code)
return logdir, metrics.finalize()
def __init__(self, logger, config, gpu):
if torch.cuda.is_available():
self.device = torch.device('cuda:{}'.format(gpu))
else:
self.device = torch.device('cpu')
random.seed(1)
numpy.random.seed(1)
torch.manual_seed(1)
torch.cuda.manual_seed_all(1)
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True
self.logger = logger
self.train_config = registry.instantiate(TrainConfig, config['train'])
self.train_config.eval_every_n = 500
self.train_config.save_every_n = 500
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(
f'Step {last_step}: loss={loss.item():.4f}')
last_step += 1
# Run saver
if last_step == 1 or last_step % self.train_config.save_every_n == 0:
saver.save(modeldir, last_step)
# Save final model
saver.save(modeldir, last_step)
def finetune(self, config, model_load_dir, model_save_dir,
infer_output_path, beam_size, output_history, use_heuristic):
random_seeds = [i for i in range(3)]
orig_data = registry.construct('dataset', self.config['data']['val'])
databases = orig_data.get_databases()
for seed in random_seeds:
data_random = random_state.RandomContext(seed)
print("seed:", seed)
metrics_list = []
batch_1_scores = []
no_grad_scores = []
batch_32_scores = []
n_2_scores = []
with data_random:
# print("No grad")
# no_grad_infer_output_path = infer_output_path + "no_grad/no_grad.infer"
# os.makedirs(os.path.dirname(no_grad_infer_output_path), exist_ok=False)
# print(no_grad_infer_output_path)
# for database in databases:
#
# self.finetune_on_database(no_grad_infer_output_path, database, config, model_load_dir,
# beam_size, output_history, use_heuristic, metrics_list, no_grad_scores,
# take_grad_steps=False, batch_size="1")
# print("No grad scores", no_grad_scores)
# print("average", self.aggregate_score(no_grad_scores))
no_grad_scores = [
('dog_kennels', 0.5, 82), ('flight_2', 0.5875, 80),
('pets_1', 0.4523809523809524, 42),
('concert_singer', 0.5333333333333333, 45),
('museum_visit', 0.4444444444444444, 18),
('battle_death', 0.5625, 16),
('student_transcripts_tracking', 0.48717948717948717, 78),
('singer', 0.7333333333333333, 30),
('cre_Doc_Template_Mgt', 0.7023809523809523, 84),
('world_1', 0.19166666666666668, 120),
('employee_hire_evaluation', 0.8421052631578947, 38),
('network_1', 0.6428571428571429, 56),
('poker_player', 0.875, 40),
('real_estate_properties', 0.25, 4),
('course_teach', 0.7333333333333333, 30),
('voter_1', 0.4666666666666667, 15), ('wta_1', 0.5, 62),
('orchestra', 0.85, 40),
('car_1', 0.32608695652173914, 92),
('tvshow', 0.6612903225806451, 62)
]
average = self.aggregate_score(no_grad_scores)
no_grad_scores.append(("average", average))
# self.plot(no_grad_scores, "no_grad_scores.png", "no grad scores")
# print("No grad scores", no_grad_scores)
# print("average", average)
#
# print("batch size 1")
# batch_1_infer_output_path = infer_output_path + "seed_"+str(seed)+"/batch_1/batch_1.infer"
# os.makedirs(os.path.dirname(batch_1_infer_output_path), exist_ok=False)
# print(batch_1_infer_output_path)
# for database in databases:
# self.finetune_on_database(batch_1_infer_output_path, database, config, model_load_dir,
# beam_size, output_history, use_heuristic, metrics_list, batch_1_scores,
# take_grad_steps=True, batch_size="1")
# average = self.aggregate_score(batch_1_scores)
# batch_1_scores.append(("average", average))
# self.plot(batch_1_scores, "batch_1_scores_seed_"+str(seed)+".png", "batch size 1 scores seed "+ str(seed))
# print("batch size 1 scores", batch_1_scores)
# print("average", average)
#
# print("batch size 32")
# batch_32_infer_output_path = infer_output_path + "seed_"+str(seed)+"/batch_32/batch_32.infer"
# os.makedirs(os.path.dirname(batch_32_infer_output_path), exist_ok=False)
# print(batch_32_infer_output_path)
# for database in databases:
# self.finetune_on_database(batch_32_infer_output_path, database, config, model_load_dir,
# beam_size, output_history, use_heuristic, metrics_list, batch_32_scores,
# take_grad_steps=True, batch_size="32")
# average = self.aggregate_score(batch_32_scores)
# batch_32_scores.append(("average", average))
# self.plot(batch_32_scores, "batch_32_scores_seed_"+str(seed)+".png", "batch size 32 scores seed " + str(seed))
# print("batch size 32 scores", batch_32_scores)
# print("average",average)
print("n^2")
n_2_infer_output_path = infer_output_path + "seed_" + str(
seed) + "/n_2/n_2.infer"
os.makedirs(os.path.dirname(n_2_infer_output_path),
exist_ok=False)
print(n_2_infer_output_path)
for database in databases:
self.finetune_on_database(n_2_infer_output_path,
database,
config,
model_load_dir,
beam_size,
output_history,
use_heuristic,
metrics_list,
n_2_scores,
take_grad_steps=True,
batch_size="n^2")
average = self.aggregate_score(n_2_scores)
n_2_scores.append(("average", average))
self.plot(n_2_scores,
"n_2_scores_no_repeat_seed_" + str(seed) + ".png",
"batch n^2 scores no repeat seed " + str(seed))
print("n^2 scores", n_2_scores)
print("average", average)
# print("Score on entire validation set:")
# self.finetune_on_database(infer_output_path, None, config, model_load_dir,
# beam_size, output_history, use_heuristic, metrics_list, scores, take_grad_steps=False)
print("")
print("changes")
# print("batch size 1 changes")
# self.plot(self.get_change(no_grad_scores, batch_1_scores),
# "batch_size_1_changes_seed_"+str(seed)+".png",
# "batch size 1 score changes")
# print(self.get_change(no_grad_scores, batch_1_scores))
#
# print("batch size 32 changes")
# self.plot(self.get_change(no_grad_scores, batch_32_scores),
# "batch_size_32_changes_seed_" + str(seed) + ".png",
# "batch size 32 score changes")
# print(self.get_change(no_grad_scores, batch_32_scores))
print("batch size n^2 changes")
self.plot(
self.get_change(no_grad_scores, n_2_scores),
"batch_size_n_2_no_repeat_changes_seed_" + str(seed) +
".png",
"batch size n^2 score changes with no repeat queries")
print(self.get_change(no_grad_scores, n_2_scores))
def finetune_on_database(self,
infer_output_path,
database,
config,
model_load_dir,
beam_size,
output_history,
use_heuristic,
metrics_list,
scores,
take_grad_steps=True,
batch_size="1"):
if database:
current_infer_output_path = infer_output_path + "/" + database
else:
current_infer_output_path = infer_output_path + "/" + "entire_val"
os.makedirs(os.path.dirname(current_infer_output_path), exist_ok=True)
infer_output = open(current_infer_output_path, 'w')
spider_data = registry.construct('dataset',
self.config['data']['val'],
database=database)
val_data = self.model_preproc.dataset('val', database=database)
# val_data_loader = self._yield_batches_from_epochs(
# torch.utils.data.DataLoader(val_data, batch_size=1, collate_fn=lambda x: x,
# shuffle=False))
assert len(val_data) == len(spider_data)
if len(val_data) == 0:
return
if batch_size == "32":
if len(val_data) < 32:
return
print("database:", database)
if batch_size == "n^2":
indices = np.random.permutation(
self.get_no_repeat_data_indices(spider_data))
print("length of data:", len(val_data))
print("length of data after removing repeat entries:",
len(indices))
else:
indices = np.random.permutation(len(val_data))
# TODO: RANDOMIZE DATA
optimizer, lr_scheduler = self.construct_optimizer_and_lr_scheduler(
config)
saver = saver_mod.Saver({
"model": self.model,
"optimizer": optimizer
},
keep_every_n=self.finetune_config.keep_every_n)
last_step = saver.restore(model_load_dir, map_location=self.device)
self.logger.log(f"Loaded trained model; last_step:{last_step}")
current_batch = []
clear_batch = False
current_number = 0
for i in tqdm.tqdm(indices):
current_number += 1
orig_item, preproc_item = spider_data[i], val_data[i]
with torch.no_grad():
decoded = self._infer_one(self.model, orig_item, preproc_item,
beam_size, output_history,
use_heuristic)
infer_output.write(
json.dumps({
'index': int(i),
'beams': decoded,
}) + '\n')
infer_output.flush()
if take_grad_steps:
if batch_size == "1":
current_batch = [preproc_item]
elif batch_size == "32":
if current_number % 32 != 0:
current_batch.append(preproc_item)
clear_batch = False
continue
else:
clear_batch = True
else:
current_batch.append(preproc_item)
try:
with self.model_random:
loss = self.model.compute_loss(current_batch)
norm_loss = loss / self.finetune_config.num_batch_accumulated
norm_loss.backward()
if self.finetune_config.clip_grad:
torch.nn.utils.clip_grad_norm_(optimizer.bert_param_group["params"], \
self.finetune_config.clip_grad)
optimizer.step()
lr_scheduler.update_lr(last_step)
optimizer.zero_grad()
if clear_batch:
current_batch = []
# stats = self._eval_model(self.logger, self.model, last_step, batch, 'val',
# self.finetune_config.report_every_n)
# val_losses.append(stats['loss'])
except KeyError:
self.logger.log("keyError")
current_batch = []
continue
# except AssertionError:
# self.logger.log("AssertionError")
# continue
inferred = open(current_infer_output_path)
metrics = spider_data.Metrics(spider_data)
inferred_lines = list(inferred)
# if len(inferred_lines) < len(spider_data):
# raise Exception(f'Not enough inferred: {len(inferred_lines)} vs {len(spider_data)}')
for line in inferred_lines:
infer_results = json.loads(line)
if infer_results['beams']:
inferred_code = infer_results['beams'][0]['inferred_code']
else:
inferred_code = None
if 'index' in infer_results:
metrics.add(spider_data[infer_results['index']], inferred_code)
else:
metrics.add(None,
inferred_code,
obsolete_gold_code=infer_results['gold_code'])
final_metrics = metrics.finalize()
metrics_list.append(final_metrics)
#print(final_metrics['total_scores']['all']['exact'])
scores.append((database, final_metrics['total_scores']['all']['exact'],
len(indices)))
# if last_step % self.finetune_config.save_every_n == 0:
# saver.save(model_save_dir+'/seed_'+seed, last_step)
#print('scores', scores)
#print("average score:", self.aggregate_score(scores))
return scores
def train(self, config, modeldir, trainset, valset):
# 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)
else:
optimizer = registry.construct('optimizer',
config['optimizer'],
params=self.model.parameters())
# 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"])
print("pretrain restored! pretrain step: %d" % pretrain_step)
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(trainset)
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(valset)
print("train: ")
for _item in train_data.components[1]:
if _item.tree is None:
print("?")
dev_badidxs = []
print("val: ")
for _idx, _item in enumerate(val_data.components[1]):
if _item.tree is None:
dev_badidxs.append(_idx)
print("!")
assert len(val_data.components[0]) == len(val_data.components[1])
new_first = []
new_second = []
for _idx in range(len(val_data.components[1])):
if _idx not in dev_badidxs:
new_first.append(val_data.components[0][_idx])
new_second.append(val_data.components[1][_idx])
val_data.components = copy.deepcopy((new_first, new_second))
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:
last_val_loss = None
lr_decay_countdown = MAX_LRDECAY_COUNTDOWN
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:
train_loss = 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:
eval_loss = self._eval_model(
self.logger,
self.model,
last_step,
val_data_loader,
'val',
num_eval_items=self.train_config.num_eval_items)
if last_val_loss is None or eval_loss < last_val_loss:
last_val_loss = eval_loss
lr_decay_countdown = MAX_LRDECAY_COUNTDOWN
elif lr_decay_countdown > 0:
lr_decay_countdown -= 1
else:
current_lr = None
for p in optimizer.param_groups:
p['lr'] *= LR_DECAY_RATE
current_lr = p['lr']
self.logger.log(f'LR decay: down to {current_lr}')
if DEBUG:
current_lr = None
for p in optimizer.param_groups:
p['lr'] *= LR_DECAY_RATE
current_lr = p['lr']
self.logger.log(f'LR decay: down to {current_lr}')
# 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()
optimizer.zero_grad()
# Report metrics
if last_step % self.train_config.report_every_n == 0:
self.logger.log(
f'Step {last_step}: loss={loss.item():.4f}')
last_step += 1
# Run saver
if last_step == 1 or last_step % self.train_config.save_every_n == 0:
saver.save(modeldir, last_step)
print("model saved at %d step" % last_step)
# Save final model
saver.save(modeldir, last_step)
