def pytorch_benchmark(batch_sizes, sequence_lengths, nums_random_blocks, output_path, attention_type="block_sparse"):
# Compare takes a list of measurements which we'll save in results.
device = torch.device("cuda")
fp = open(output_path, "w")
writer = csv.writer(fp)
writer.writerow(["batch_size", "seq_length", "r", "forward time (ms)", "bakward time (ms)"])
tokenizer = BigBirdTokenizer.from_pretrained("google/bigbird-roberta-base")
for b, n, r in product(batch_sizes, sequence_lengths, nums_random_blocks):
print(b, n, r)
inputs = tokenizer([input_text for _ in range(b)], max_length=n, truncation=True, return_tensors="pt")
config = BigBirdConfig.from_pretrained("google/bigbird-roberta-base", attention_type=attention_type)
model = BigBirdForSequenceClassification.from_pretrained("google/bigbird-roberta-base", config=config)
model.to(device)
try:
torch.cuda.synchronize()
forward_time = 0
backward_time = 0
for _ in range(10):
forward_elapse, backward_elapse = time_foward_backward(model, inputs)
forward_time += forward_elapse
backward_time += backward_elapse
forward_time /= 10
backward_time /= 10
print(forward_time, backward_time)
writer.writerow([b, n, r, forward_time, backward_time])
except Exception as e:
print("Error:", e)
traceback.print_exc()
fp.close()
def create_and_check_for_sequence_classification(
self, config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
):
config.num_labels = self.num_labels
model = BigBirdForSequenceClassification(config)
model.to(torch_device)
model.eval()
result = model(input_ids, attention_mask=input_mask, token_type_ids=token_type_ids, labels=sequence_labels)
self.parent.assertEqual(result.logits.shape, (self.batch_size, self.num_labels))