def create_and_check_reformer_random_seed(self, config, input_ids,
input_mask, choice_labels):
layer = ReformerLayer(config).to(torch_device)
layer.train()
shape = (
self.batch_size,
self.seq_length,
config.hidden_size,
) # Batch x SeqLen x hiddenSize
hidden_states = floats_tensor(shape)
attn_output = floats_tensor(shape)
seeds = []
for _ in range(100):
layer_outputs = layer(attn_output,
hidden_states,
attention_mask=input_mask)
attn_output = layer_outputs.attn_output
hidden_states = layer_outputs.hidden_states
torch.manual_seed(layer.attention_seed)
seeds.append(layer.attention_seed)
self.parent.assertGreater(len(set(seeds)), 70)
seeds = []
for _ in range(100):
layer_outputs = layer(attn_output,
hidden_states,
attention_mask=input_mask)
attn_output = layer_outputs.attn_output
hidden_states = layer_outputs.hidden_states
torch.manual_seed(layer.feed_forward_seed)
seeds.append(layer.feed_forward_seed)
self.parent.assertGreater(len(set(seeds)), 70)
def test_local_layer_forward_complex(self):
config = self._get_basic_config_and_input()
config["attn_layers"] = ["local"]
attn_mask = self._get_attn_mask()
hidden_states = self._get_hidden_states()
torch.manual_seed(0)
layer = ReformerLayer(ReformerConfig(**config)).to(torch_device)
layer.eval()
reformer_output = layer(prev_attn_output=hidden_states, hidden_states=hidden_states, attention_mask=attn_mask,)
output_slice = reformer_output.hidden_states[0, 0, :5]
expected_output_slice = torch.tensor(
[1.5476, -1.9020, -0.9902, 1.5013, -0.1950], dtype=torch.float, device=torch_device,
)
self.assertTrue(torch.allclose(output_slice, expected_output_slice, atol=1e-3))
def test_local_layer_forward(self):
config = self._get_basic_config_and_input()
config["attn_layers"] = ["local"]
config["is_decoder"] = False
hidden_states = self._get_hidden_states()
torch.manual_seed(0)
layer = ReformerLayer(ReformerConfig(**config)).to(torch_device)
layer.eval()
reformer_output = layer(prev_attn_output=hidden_states, hidden_states=hidden_states)
output_slice = reformer_output.hidden_states[0, 0, :5]
expected_output_slice = torch.tensor(
[1.4212, -2.0576, -0.9688, 1.4599, -0.1344], dtype=torch.float, device=torch_device,
)
self.assertTrue(torch.allclose(output_slice, expected_output_slice, atol=1e-3))
def test_lsh_layer_forward(self):
config = self._get_basic_config_and_input()
config["attn_layers"] = ["lsh"]
config["is_decoder"] = False
hidden_states = self._get_hidden_states()
torch.manual_seed(0)
layer = ReformerLayer(ReformerConfig(**config)).to(torch_device)
layer.eval()
reformer_output = layer(prev_attn_output=hidden_states.clone(), hidden_states=hidden_states)
output_slice = reformer_output.hidden_states[0, 0, :5]
expected_output_slice = torch.tensor(
[1.6879, -1.3083, -0.4708, 1.3555, -0.6292], dtype=torch.float, device=torch_device,
)
self.assertTrue(torch.allclose(output_slice, expected_output_slice, atol=1e-3))
def test_lsh_layer_forward_complex(self):
config = self._get_basic_config_and_input()
config["attn_layers"] = ["lsh"]
config["num_buckets"] = [2, 4]
attn_mask = self._get_attn_mask()
hidden_states = self._get_hidden_states()
torch.manual_seed(0)
layer = ReformerLayer(ReformerConfig(**config)).to(torch_device)
layer.eval()
reformer_output = layer(
prev_attn_output=hidden_states.clone(), hidden_states=hidden_states, attention_mask=attn_mask,
)
output_slice = reformer_output.hidden_states[0, 0, :5]
expected_output_slice = torch.tensor(
[1.6439, -1.2306, -0.5108, 1.3006, -0.6537], dtype=torch.float, device=torch_device,
)
self.assertTrue(torch.allclose(output_slice, expected_output_slice, atol=1e-3))
def create_and_check_reformer_layer_dropout_seed(self, config, input_ids, input_mask, is_decoder):
config.is_decoder = is_decoder
layer = ReformerLayer(config).to(torch_device)
layer.train()
shape = (
self.batch_size,
self.seq_length,
config.hidden_size,
) # Batch x SeqLen x hiddenSize
# get random tensors
hidden_states = floats_tensor(shape)
prev_attn_output = floats_tensor(shape)
# now the random seeds for attention and feed forward is initialized
# forward tensors with dropout
layer_outputs = layer(prev_attn_output, hidden_states, attention_mask=input_mask)
next_attn_output = layer_outputs.attn_output
next_hidden_states = layer_outputs.hidden_states
torch.manual_seed(layer.attention_seed)
attn_outputs = layer.attention(hidden_states, attention_mask=input_mask)
self.parent.assertTrue(
torch.allclose(prev_attn_output + attn_outputs.hidden_states, next_attn_output, atol=1e-3,)
)
torch.manual_seed(layer.feed_forward_seed)
feed_forward_hidden_states = layer.feed_forward(next_attn_output)
self.parent.assertTrue(
torch.allclose(next_hidden_states, hidden_states + feed_forward_hidden_states, atol=1e-3,)
)
def _compute_pytorch(
model_names,
batch_sizes,
slice_sizes,
dictionary,
average_over,
device,
torchscript,
fp16,
no_speed,
no_memory,
verbose,
num_hashes
):
hidden_size = 64
num_attention_heads = 2
intermediate_size = 128
chunk_length = 64
num_hashes = num_hashes
hidden_states = floats_tensor((1, 2 ** 16, hidden_size))
for c, model_name in enumerate(model_names):
print(f"{c + 1} / {len(model_names)}")
dictionary[model_name] = {
"bs": batch_sizes,
"ss": slice_sizes,
"results": {},
"memory": {},
}
dictionary[model_name]["results"] = {i: {} for i in batch_sizes}
dictionary[model_name]["memory"] = {i: {} for i in batch_sizes}
for batch_size in batch_sizes:
for slice_size in slice_sizes:
num_buckets = int(2 * slice_size / chunk_length)
if num_buckets > chunk_length:
factorized_num_buckets = num_buckets // 32
num_buckets = [32, factorized_num_buckets]
bert_config = BertConfig(
hidden_size=hidden_size,
num_attention_heads=num_attention_heads,
intermediate_size=intermediate_size,
hidden_dropout_prob=0.0,
attention_probs_dropout_prob=0.0,
)
reformer_config = ReformerConfig(
hidden_size=hidden_size,
num_attention_heads=num_attention_heads,
intermediate_size=intermediate_size,
chunk_length=chunk_length,
num_hashes=num_hashes,
num_buckets=num_buckets
)
layers = {
'ReformerLayer': ReformerLayer(reformer_config),
'BertLayer': BertLayer(bert_config)
}
model = layers[model_name]
if fp16:
model.half()
model.to(device)
model.eval()
if False:
dictionary[model_name]["results"][batch_size][slice_size] = "N/A"
else:
sequence = (
hidden_states[0, :slice_size, :]
.to(device=device)
.repeat(batch_size, 1, 1)
)
try:
if torchscript:
print("Tracing model with sequence size", sequence.shape)
inference = torch.jit.trace(model, sequence)
inference(sequence)
else:
inference = model
if model_name == "ReformerLayer":
inference(sequence, sequence)
else:
inference(sequence)
if not no_memory:
# model.add_memory_hooks() # Forward method tracing (only for PyTorch models)
trace = start_memory_tracing("transformers")
if model_name == "ReformerLayer":
inference(sequence, sequence)
else:
inference(sequence)
summary = stop_memory_tracing(trace)
if verbose:
print_summary_statistics(summary)
dictionary[model_name]["memory"][batch_size][
slice_size
] = str(summary.total)
else:
dictionary[model_name]["memory"][batch_size][
slice_size
] = "N/A"
if not no_speed:
print(
"Going through model with sequence of shape",
sequence.shape,
)
if model_name == "ReformerLayer":
runtimes = timeit.repeat(
lambda: inference(sequence, sequence),
repeat=average_over,
number=3,
)
else:
runtimes = timeit.repeat(
lambda: inference(sequence),
repeat=average_over,
number=3,
)
average_time = sum(runtimes) / float(len(runtimes)) / 3.0
dictionary[model_name]["results"][batch_size][
slice_size
] = average_time
else:
dictionary[model_name]["results"][batch_size][
slice_size
] = "N/A"
except RuntimeError as e:
print("Doesn't fit on GPU.", e)
torch.cuda.empty_cache()
dictionary[model_name]["results"][batch_size][
slice_size
] = "N/A"
dictionary[model_name]["memory"][batch_size][slice_size] = "N/A"
return dictionary
