def check_save_and_load(self, config, input_ids, attention_mask,
encoder_hidden_states, decoder_config,
decoder_input_ids, decoder_attention_mask,
**kwargs):
encoder_model, decoder_model = self.get_encoder_decoder_model(
config, decoder_config)
enc_dec_model = EncoderDecoderModel(encoder=encoder_model,
decoder=decoder_model)
enc_dec_model.to(torch_device)
enc_dec_model.eval()
with torch.no_grad():
outputs = enc_dec_model(
input_ids=input_ids,
decoder_input_ids=decoder_input_ids,
attention_mask=attention_mask,
decoder_attention_mask=decoder_attention_mask,
)
out_2 = outputs[0].cpu().numpy()
out_2[np.isnan(out_2)] = 0
with tempfile.TemporaryDirectory() as tmpdirname:
enc_dec_model.save_pretrained(tmpdirname)
enc_dec_model = EncoderDecoderModel.from_pretrained(tmpdirname)
enc_dec_model.to(torch_device)
after_outputs = enc_dec_model(
input_ids=input_ids,
decoder_input_ids=decoder_input_ids,
attention_mask=attention_mask,
decoder_attention_mask=decoder_attention_mask,
)
out_1 = after_outputs[0].cpu().numpy()
out_1[np.isnan(out_1)] = 0
max_diff = np.amax(np.abs(out_1 - out_2))
self.assertLessEqual(max_diff, 1e-5)
def check_encoder_decoder_model_from_pretrained_configs(
self,
config,
input_ids,
attention_mask,
encoder_hidden_states,
decoder_config,
decoder_input_ids,
decoder_attention_mask,
**kwargs
):
encoder_decoder_config = EncoderDecoderConfig.from_encoder_decoder_configs(config, decoder_config)
self.assertTrue(encoder_decoder_config.decoder.is_decoder)
enc_dec_model = EncoderDecoderModel(encoder_decoder_config)
enc_dec_model.to(torch_device)
enc_dec_model.eval()
self.assertTrue(enc_dec_model.config.is_encoder_decoder)
outputs_encoder_decoder = enc_dec_model(
input_ids=input_ids,
decoder_input_ids=decoder_input_ids,
attention_mask=attention_mask,
decoder_attention_mask=decoder_attention_mask,
)
self.assertEqual(outputs_encoder_decoder[0].shape, (decoder_input_ids.shape + (decoder_config.vocab_size,)))
self.assertEqual(outputs_encoder_decoder[1].shape, (input_ids.shape + (config.hidden_size,)))
def create_and_check_encoder_decoder_shared_weights(
self, config, input_ids, attention_mask, encoder_hidden_states,
decoder_config, decoder_input_ids, decoder_attention_mask, labels,
**kwargs):
torch.manual_seed(0)
encoder_model, decoder_model = self.get_encoder_decoder_model(
config, decoder_config)
model = EncoderDecoderModel(encoder=encoder_model,
decoder=decoder_model)
model.to(torch_device)
model.eval()
# load state dict copies weights but does not tie them
decoder_state_dict = model.decoder._modules[
model.decoder.base_model_prefix].state_dict()
model.encoder.load_state_dict(decoder_state_dict, strict=False)
torch.manual_seed(0)
tied_encoder_model, tied_decoder_model = self.get_encoder_decoder_model(
config, decoder_config)
config = EncoderDecoderConfig.from_encoder_decoder_configs(
tied_encoder_model.config,
tied_decoder_model.config,
tie_encoder_decoder=True)
tied_model = EncoderDecoderModel(encoder=tied_encoder_model,
decoder=tied_decoder_model,
config=config)
tied_model.to(torch_device)
tied_model.eval()
model_result = model(
input_ids=input_ids,
decoder_input_ids=decoder_input_ids,
attention_mask=attention_mask,
decoder_attention_mask=decoder_attention_mask,
)
tied_model_result = tied_model(
input_ids=input_ids,
decoder_input_ids=decoder_input_ids,
attention_mask=attention_mask,
decoder_attention_mask=decoder_attention_mask,
)
# check that models has less parameters
self.assertLess(sum(p.numel() for p in tied_model.parameters()),
sum(p.numel() for p in model.parameters()))
random_slice_idx = ids_tensor((1, ), model_result[0].shape[-1]).item()
# check that outputs are equal
self.assertTrue(
torch.allclose(model_result[0][0, :, random_slice_idx],
tied_model_result[0][0, :, random_slice_idx],
atol=1e-4))
# check that outputs after saving and loading are equal
with tempfile.TemporaryDirectory() as tmpdirname:
tied_model.save_pretrained(tmpdirname)
tied_model = EncoderDecoderModel.from_pretrained(tmpdirname)
tied_model.to(torch_device)
tied_model.eval()
# check that models has less parameters
self.assertLess(sum(p.numel() for p in tied_model.parameters()),
sum(p.numel() for p in model.parameters()))
random_slice_idx = ids_tensor((1, ),
model_result[0].shape[-1]).item()
tied_model_result = tied_model(
input_ids=input_ids,
decoder_input_ids=decoder_input_ids,
attention_mask=attention_mask,
decoder_attention_mask=decoder_attention_mask,
)
# check that outputs are equal
self.assertTrue(
torch.allclose(model_result[0][0, :, random_slice_idx],
tied_model_result[0][0, :, random_slice_idx],
atol=1e-4))
def sample_generate(top_k=50,
temperature=1.0,
model_path='/content/BERT checkpoints/model-9.pth',
gpu_id=0):
# make sure your model is on GPU
device = torch.device(f"cuda:{gpu_id}")
# ------------------------LOAD MODEL-----------------
print('load the model....')
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
bert_encoder = BertConfig.from_pretrained('bert-base-uncased')
bert_decoder = BertConfig.from_pretrained('bert-base-uncased',
is_decoder=True)
config = EncoderDecoderConfig.from_encoder_decoder_configs(
bert_encoder, bert_decoder)
model = EncoderDecoderModel(config)
model.load_state_dict(torch.load(model_path, map_location='cuda'))
model = model.to(device)
encoder = model.get_encoder()
decoder = model.get_decoder()
model.eval()
print('load success')
# ------------------------END LOAD MODEL--------------
# ------------------------LOAD VALIDATE DATA------------------
test_data = torch.load("/content/test_data.pth")
test_dataset = TensorDataset(*test_data)
test_dataloader = DataLoader(dataset=test_dataset,
shuffle=False,
batch_size=1)
# ------------------------END LOAD VALIDATE DATA--------------
# ------------------------START GENERETE-------------------
update_count = 0
bleu_2scores = 0
bleu_4scores = 0
nist_2scores = 0
nist_4scores = 0
sentences = []
meteor_scores = 0
print('start generating....')
for batch in test_dataloader:
with torch.no_grad():
batch = [item.to(device) for item in batch]
encoder_input, decoder_input, mask_encoder_input, _ = batch
past, _ = encoder(encoder_input, mask_encoder_input)
prev_pred = decoder_input[:, :1]
sentence = prev_pred
# decoding loop
for i in range(100):
logits = decoder(sentence, encoder_hidden_states=past)
logits = logits[0][:, -1]
logits = logits.squeeze(1) / temperature
logits = top_k_logits(logits, k=top_k)
probs = F.softmax(logits, dim=-1)
prev_pred = torch.multinomial(probs, num_samples=1)
sentence = torch.cat([sentence, prev_pred], dim=-1)
if prev_pred[0][0] == 102:
break
predict = tokenizer.convert_ids_to_tokens(sentence[0].tolist())
encoder_input = encoder_input.squeeze(dim=0)
encoder_input_num = (encoder_input != 0).sum()
inputs = tokenizer.convert_ids_to_tokens(
encoder_input[:encoder_input_num].tolist())
decoder_input = decoder_input.squeeze(dim=0)
decoder_input_num = (decoder_input != 0).sum()
reference = tokenizer.convert_ids_to_tokens(
decoder_input[:decoder_input_num].tolist())
print('-' * 20 + f"example {update_count}" + '-' * 20)
print(f"input: {' '.join(inputs)}")
print(f"output: {' '.join(reference)}")
print(f"predict: {' '.join(predict)}")
temp_bleu_2, \
temp_bleu_4, \
temp_nist_2, \
temp_nist_4, \
temp_meteor_scores = calculate_metrics(predict[1:-1], reference[1:-1])
bleu_2scores += temp_bleu_2
bleu_4scores += temp_bleu_4
nist_2scores += temp_nist_2
nist_4scores += temp_nist_4
meteor_scores += temp_meteor_scores
sentences.append(" ".join(predict[1:-1]))
update_count += 1
entro, dist = cal_entropy(sentences)
mean_len, var_len = cal_length(sentences)
print(f'avg: {mean_len}, var: {var_len}')
print(f'entro: {entro}')
print(f'dist: {dist}')
print(f'test bleu_2scores: {bleu_2scores / update_count}')
print(f'test bleu_4scores: {bleu_4scores / update_count}')
print(f'test nist_2scores: {nist_2scores / update_count}')
print(f'test nist_4scores: {nist_4scores / update_count}')
print(f'test meteor_scores: {meteor_scores / update_count}')
def train_model(epochs=10,
num_gradients_accumulation=4,
batch_size=4,
gpu_id=0,
lr=1e-5,
load_dir='/content/BERT checkpoints'):
# make sure your model is on GPU
device = torch.device(f"cuda:{gpu_id}")
# ------------------------LOAD MODEL-----------------
print('load the model....')
bert_encoder = BertConfig.from_pretrained('bert-base-uncased')
bert_decoder = BertConfig.from_pretrained('bert-base-uncased',
is_decoder=True)
config = EncoderDecoderConfig.from_encoder_decoder_configs(
bert_encoder, bert_decoder)
model = EncoderDecoderModel(config)
model = model.to(device)
print('load success')
# ------------------------END LOAD MODEL--------------
# ------------------------LOAD TRAIN DATA------------------
train_data = torch.load("/content/train_data.pth")
train_dataset = TensorDataset(*train_data)
train_dataloader = DataLoader(dataset=train_dataset,
shuffle=True,
batch_size=batch_size)
val_data = torch.load("/content/validate_data.pth")
val_dataset = TensorDataset(*val_data)
val_dataloader = DataLoader(dataset=val_dataset,
shuffle=True,
batch_size=batch_size)
# ------------------------END LOAD TRAIN DATA--------------
# ------------------------SET OPTIMIZER-------------------
num_train_optimization_steps = len(
train_dataset) * epochs // batch_size // num_gradients_accumulation
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,
weight_decay=0.01,
)
scheduler = get_linear_schedule_with_warmup(
optimizer,
num_warmup_steps=num_train_optimization_steps // 10,
num_training_steps=num_train_optimization_steps)
# ------------------------START TRAINING-------------------
update_count = 0
start = time.time()
print('start training....')
for epoch in range(epochs):
# ------------------------training------------------------
model.train()
losses = 0
times = 0
print('\n' + '-' * 20 + f'epoch {epoch}' + '-' * 20)
for batch in tqdm(train_dataloader):
batch = [item.to(device) for item in batch]
encoder_input, decoder_input, mask_encoder_input, mask_decoder_input = batch
logits = model(input_ids=encoder_input,
attention_mask=mask_encoder_input,
decoder_input_ids=decoder_input,
decoder_attention_mask=mask_decoder_input)
out = logits[0][:, :-1].contiguous()
target = decoder_input[:, 1:].contiguous()
target_mask = mask_decoder_input[:, 1:].contiguous()
loss = util.sequence_cross_entropy_with_logits(out,
target,
target_mask,
average="token")
loss.backward()
losses += loss.item()
times += 1
update_count += 1
if update_count % num_gradients_accumulation == num_gradients_accumulation - 1:
torch.nn.utils.clip_grad_norm_(model.parameters(),
max_grad_norm)
optimizer.step()
scheduler.step()
optimizer.zero_grad()
end = time.time()
print(f'time: {(end - start)}')
print(f'loss: {losses / times}')
start = end
# ------------------------validate------------------------
model.eval()
perplexity = 0
batch_count = 0
print('\nstart calculate the perplexity....')
with torch.no_grad():
for batch in tqdm(val_dataloader):
batch = [item.to(device) for item in batch]
encoder_input, decoder_input, mask_encoder_input, mask_decoder_input = batch
logits = model(input_ids=encoder_input,
attention_mask=mask_encoder_input,
decoder_input_ids=decoder_input,
decoder_attention_mask=mask_decoder_input)
out = logits[0][:, :-1].contiguous()
target = decoder_input[:, 1:].contiguous()
target_mask = mask_decoder_input[:, 1:].contiguous()
# print(out.shape,target.shape,target_mask.shape)
loss = util.sequence_cross_entropy_with_logits(out,
target,
target_mask,
average="token")
perplexity += np.exp(loss.item())
batch_count += 1
print(f'\nvalidate perplexity: {perplexity / batch_count}')
torch.save(
model.state_dict(),
os.path.join(os.path.abspath('.'), load_dir,
"model-" + str(epoch) + ".pth"))
optimizer.zero_grad()
train_loss.backward()
optimizer.step()
train_acc = calculate_accuracy(outputs.logits, labels)
train_loss_list.append(train_loss.item())
train_acc_list.append(train_acc)
mean_train_loss = sum(train_loss_list) / (len(train_loss_list) + 1e-4)
mean_train_acc = sum(train_acc_list) / (len(train_acc_list) + 1e-4)
print("epoch: {} train_loss: {:.3f}, train_acc: {:.3f}".format(epoch, mean_train_loss, mean_train_acc))
if epoch % 5 == 0:
model.eval()
valid_all_match = []
for tasks, plans in tqdm(valid_seen_dataloader):
try:
tokenized_text = encoder_tokenizer(tasks, padding=True, truncation=True, max_length=100, return_tensors="pt").input_ids
if args.gpu and torch.cuda.is_available():
tokenized_text = tokenized_text.to("cuda")
output_labels = model.generate(tokenized_text, decoder_start_token_id=1)
ouput_array = output_labels.cpu().numpy()[0]
targets = decoder_tokenizer.tokenize(plans)
all_match = np.all(ouput_array[:len(targets[0])] == targets[0])
valid_all_match.append(1 if all_match else 0)
