c-my/transformers
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

make lm model work

Browse Source
This commit is contained in:
Patrick von Platen committed 2020-04-22 07:55:19 +02:00
1 parent 062046c4eb
commit d4c347dfd5
3 files changed
+141 -1140
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/transformers/@!
-382
View File
@@ -1,382 +0,0 @@
# coding=utf-8 # Copyright 2020 Huggingface
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import unittest
import numpy as np
# trax imports - to be deleted later
from trax import math as trax_math
from trax.shapes import ShapeDtype as trax_ShapeDtype
import jax
from trax.layers.research.efficient_attention_v2 import (
LSHSelfAttention as TraxLSHSelfAttention,
)
from trax.models.reformer.reformer import DecoderBlock as TraxLSHAttentionBlock
from trax.models.reformer.reformer import ReformerLM as TraxReformer
from trax import layers as tl
from transformers import ReformerAttention, ReformerLayer, ReformerConfig
from transformers import is_torch_available # noqa: F401
from .utils import require_torch, torch_device # noqa: F401
if is_torch_available():
import torch # noqa: F401
# from transformers.modeling_reformer import ()
PATH_TO_SAVE_WEIGHTS = "/home/patrick/hugging_face/experiments/reformer/intermediate_weights"
class TraxUtils(object):
""" class that will help for testing in the beginning
should be deleted step-by-step
README (HOW-TO-INSTALL TRAX):
1) git clone https://github.com/patrickvonplaten/trax.git
- I had to do one tiny change to make the imports work,
see: https://github.com/patrickvonplaten/trax/commit/6c23e88afe7f1c57b0c38eeaa4d450e5f912590c)
2) link your PYTHON_PATH to ~/trax/trax
3) pip install all the missing packages HINT: the package gin is installed
- HINT: the package gin is installed with pip install gin-config==0.1.4
and not pip install gin.
- The other packages can just be installed with pip install <package> form
error message "<package> missing"
"""
def __init__(self, shape):
self._shape = shape
def convert_to_jax_array(self, np_array):
return jax.numpy.asarray(np_array)
def get_input_signature(self, shape=None, dtype=trax_math.numpy.float32):
with trax_math.use_backend("jax"):
if shape is None:
shape = self._shape
input_signature = trax_ShapeDtype(shape, dtype)
return input_signature
def get_layer(
self,
config,
use_reference_code=True,
mode="eval",
path_to_save_weights="/home/patrick/hugging_face/experiments/reformer/intermediate_weights",
**kwargs
):
with trax_math.use_backend("jax"):
hidden_size_per_head = config.hidden_size // config.num_attention_heads
layer = TraxLSHSelfAttention(
n_heads=config.num_attention_heads,
d_qk=hidden_size_per_head,
d_v=hidden_size_per_head,
chunk_len=config.chunk_length,
n_chunks_before=config.num_chunks_before,
n_chunks_after=config.num_chunks_after,
n_hashes=config.num_hashes,
n_buckets=config.num_buckets,
attention_dropout=config.attention_probs_dropout_prob,
output_dropout=config.hidden_dropout_prob,
hash_seed=config.seed,
causal=config.is_decoder,
use_reference_code=use_reference_code,
mode=mode,
path_to_save_weights=path_to_save_weights
)
return layer
def forward_layer(
self,
np_input_data,
layer,
input_signature=None,
random_number_generator=None,
):
with trax_math.use_backend("jax"):
input_data = self.convert_to_jax_array(np_input_data)
if input_signature is None:
input_signature = self.get_input_signature()
weights, state = layer.init(input_signature)
if random_number_generator is None:
random_number_generator = layer.new_rngs(1)[0]
output = layer(
input_data, weights=weights, state=state, rng=random_number_generator
)
return output, weights, state
def get_block(
self,
config,
use_reference_code=True,
share_qk=True,
ff_use_sru=0,
mode="eval",
path_to_save_weights=PATH_TO_SAVE_WEIGHTS,
):
with trax_math.use_backend("jax"):
with jax.disable_jit():
hidden_size_per_head = config.hidden_size // config.num_attention_heads
list_of_layers = TraxLSHAttentionBlock(
d_model=config.d_model,
d_ff=config.d_ff,
d_attention_key=hidden_size_per_head,
d_attention_value=hidden_size_per_head,
n_heads=config.num_attention_heads,
n_attention_chunks=config.num_attention_chunks,
attention_type=tl.LSHSelfAttention,
dropout=config.hidden_dropout_prob,
share_qk=share_qk,
ff_activation=tl.Gelu,
ff_use_sru=ff_use_sru,
ff_chunk_size=config.ff_chunk_size,
mode=mode,
causal=config.is_decoder,
chunk_len=config.chunk_length,
n_chunks_before=config.num_chunks_before,
n_chunks_after=config.num_chunks_after,
n_hashes=config.num_hashes,
n_buckets=config.num_buckets,
use_reference_code=use_reference_code,
hash_seed=config.seed,
path_to_save_weights=path_to_save_weights
)
block = tl.Serial(tl.ReversibleSerial([list_of_layers]))
return block
def forward_block(
self,
np_input_data,
block,
input_signature=None,
random_number_generator=None,
):
with trax_math.use_backend("jax"):
input_data = self.convert_to_jax_array(np_input_data)
input_data = (input_data,) * 2
if input_signature is None:
input_signature = self.get_input_signature()
input_signature = (input_signature, input_signature)
weights, state = block.init(input_signature)
if random_number_generator is None:
random_number_generator = block.new_rngs(1)[0]
output = block(
input_data, weights=weights, state=state, rng=random_number_generator
)
return output, weights, state
def get_model(
self,
config,
use_reference_code=True,
share_qk=True,
ff_use_sru=0,
axial_pos_shape=(),
d_axial_pos_embs=None,
mode="eval",
path_to_save_weights=PATH_TO_SAVE_WEIGHTS,
):
with trax_math.use_backend("jax"):
with jax.disable_jit():
hidden_size_per_head = config.hidden_size // config.num_attention_heads
model = TraxReformer(
vocab_size=config.vocab_size,
d_model=config.d_model,
d_ff=config.d_ff,
d_attention_key=hidden_size_per_head,
d_attention_value=hidden_size_per_head,
n_layers=config.num_hidden_layers,
n_heads=config.num_attention_heads,
dropout=config.hidden_dropout_prob,
max_len=config.max_position_embeddings,
n_chunks=config.num_chunks,
n_attention_chunks=config.num_attention_chunks,
attention_type=tl.LSHSelfAttention,
share_qk=share_qk,
axial_pos_shape=axial_pos_shape,
d_axial_pos_embs=d_axial_pos_embs,
ff_activation=tl.Gelu,
ff_use_sru=ff_use_sru,
ff_chunk_size=config.ff_chunk_size,
mode=mode,
causal=config.is_decoder,
chunk_len=config.chunk_length,
n_chunks_before=config.num_chunks_before,
n_chunks_after=config.num_chunks_after,
n_hashes=config.num_hashes,
n_buckets=config.num_buckets,
use_reference_code=use_reference_code,
hash_seed=config.seed,
path_to_save_weights=path_to_save_weights
)
return model
def forward_model(
self,
np_input_data,
model,
input_signature=None,
random_number_generator=None,
):
with trax_math.use_backend("jax"):
input_data = self.convert_to_jax_array(np_input_data)
input_data = (input_data,) * 2
if input_signature is None:
input_signature = self.get_input_signature(dtype=trax_math.numpy.int32)
input_signature = (input_signature, input_signature)
weights, state = model.init(input_signature)
if random_number_generator is None:
random_number_generator = model.new_rngs(1)[0]
output = model(
input_data, weights=weights, state=state, rng=random_number_generator
)
return output, weights, state
@require_torch
class ReformerIntegrationTests(unittest.TestCase):
def _set_param(self, torch_layer, weight, bias=None):
with torch.no_grad():
assert torch_layer.weight.shape == weight.shape, "{} layer.weight does not match".format(torch.layer)
torch_layer.weight = torch.nn.Parameter(weight)
if bias is not None:
assert torch_layer.bias.shape == bias.shape, "{} layer.bias does not match".format(torch.layer)
torch_layer.bias = torch.nn.Parameter(bias)
def _set_layer_weights_in_torch(self, weights, torch_layer, hidden_size):
# set torch weights for 1-to-1 comparison
np_query_key = np.asarray(weights[0])
np_value = np.asarray(weights[1])
np_dense = np.asarray(weights[2])
self._set_param(torch_layer.self_attention.query_key, torch.tensor(np_query_key).transpose(1, 2).contiguous().view(-1, hidden_size))
self._set_param(torch_layer.self_attention.value, torch.tensor(np_value).transpose(1, 2).contiguous().view(-1, hidden_size))
self._set_param(torch_layer.output.dense, torch.tensor(np_dense).view(-1, hidden_size).contiguous().transpose(0, 1))
def _set_block_weights_in_torch(self, weights, torch_layer, hidden_size):
weights = weights[0]
# layernorm 1
layer_norm_1 = weights[0][0][0]
layer_norm_1_weight = np.asarray(layer_norm_1[0])
layer_norm_1_bias = np.asarray(layer_norm_1[1])
self._set_param(torch_layer.attention.layer_norm, torch.tensor(layer_norm_1_weight), torch.tensor(layer_norm_1_bias))
# lsh weights + output
lsh_weights = weights[0][1]
self._set_layer_weights_in_torch(lsh_weights, torch_layer.attention, hidden_size)
# intermediate weighs
intermediate_weights = weights[2][0][2][2]
# layernorm 2
layer_norm_2_weight = np.asarray(intermediate_weights[0][0])
layer_norm_2_bias = np.asarray(intermediate_weights[0][1])
self._set_param(torch_layer.feed_forward.layer_norm, torch.tensor(layer_norm_2_weight), torch.tensor(layer_norm_2_bias))
# intermediate dense
inter_dense_weight = np.asarray(intermediate_weights[1][0])
inter_dense_bias = np.asarray(intermediate_weights[1][1])
self._set_param(torch_layer.feed_forward.dense.dense, torch.tensor(inter_dense_weight).transpose(0, 1).contiguous(), torch.tensor(inter_dense_bias))
# intermediate out
out_dense_weight = np.asarray(intermediate_weights[4][0])
out_dense_bias = np.asarray(intermediate_weights[4][1])
self._set_param(torch_layer.feed_forward.output.dense, torch.tensor(out_dense_weight).transpose(0, 1).contiguous(), torch.tensor(out_dense_bias))
def test_lsh_layer(self):
# Remove residual connection in ReformerSelfOutput to test this layer only
# Remove layer norm in ReformerAttention to test this layer only
config = ReformerConfig()
shape = (2, 7, config.hidden_size) # Batch x SeqLen x hiddenSize
np_input = np.random.rand(*shape)
trax_utils = TraxUtils(shape)
trax_layer = trax_utils.get_layer(config)
trax_output, trax_weights, trax_state = trax_utils.forward_layer(np_input, layer=trax_layer)
hf_input = torch.tensor(np_input, dtype=torch.float)
hf_layer = ReformerAttention(config)
self._set_layer_weights_in_torch(trax_weights, hf_layer, config.hidden_size)
hf_layer.eval()
hf_attention_all_heads = hf_layer.self_attention(hf_input)[0]
hf_output = hf_layer.output(hf_attention_all_heads, torch.zeros_like(hf_input))
trax_torch_output = torch.tensor(np.asarray(trax_output))
self.assertTrue(torch.allclose(hf_output, trax_torch_output, atol=1e-3))
def test_lsh_block(self):
config = ReformerConfig()
shape = (2, 7, config.hidden_size) # Batch x SeqLen x ModelDimPerHead
np_input = np.random.rand(*shape)
trax_utils = TraxUtils(shape)
trax_block = trax_utils.get_block(config)
trax_output, trax_weights, trax_state = trax_utils.forward_block(np_input, block=trax_block)
trax_torch_output_1 = torch.tensor(np.asarray(trax_output[0]))
trax_torch_output_2 = torch.tensor(np.asarray(trax_output[1]))
hf_input = torch.tensor(np_input, dtype=torch.float)
hf_block = ReformerLayer(config)
self._set_block_weights_in_torch(trax_weights, hf_block, config.hidden_size)
hf_block.eval()
hf_output_1, hf_output_2 = hf_block(hf_input, hf_input)[:2]
self.assertTrue(torch.allclose(hf_output_1, trax_torch_output_1, atol=1e-3))
self.assertTrue(torch.allclose(hf_output_2, trax_torch_output_2, atol=1e-3))
def test_reformer_model(self):
config = ReformerConfig()
shape = (2, 7) # Batch x SeqLen x ModelDimPerHead
np_input = np.random.randint(0, config.vocab_size, size=shape)
np_zeros = np.zeros((shape[0], 1), dtype=np.int32)
import ipdb
ipdb.set_trace()
torch_np_input = torch.cat([torch.tensor(np_zeros), torch.tensor(np_input[:, :-1])], dim=-1)
trax_utils = TraxUtils(shape)
trax_model = trax_utils.get_model(config)
trax_output, trax_weights, trax_state = trax_utils.forward_model(np_input, model=trax_model)
import ipdb
ipdb.set_trace()
src/transformers/w!
-744
View File
@@ -1,744 +0,0 @@
# coding=utf-8
# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""PyTorch REFORMER model. """
import logging
import torch
import numpy as np
from torch import nn
from .activations import gelu, gelu_new, swish
from .configuration_reformer import ReformerConfig
from .modeling_utils import PreTrainedModel
# DELETE later
import numpy
logger = logging.getLogger(__name__)
REFORMER_PRETRAINED_MODEL_ARCHIVE_MAP = {}
# TODO: fill with pretrained model weights
# "reformer-base-uncased": "https://s3.amazonaws.com/models.huggingface.co/reformer/reformer-base-uncased-pytorch_model.bin",
# "reformer-large-uncased": "https://s3.amazonaws.com/models.huggingface.co/reformer/reformer-large-uncased-pytorch_model.bin",
# "reformer-base-cased": "https://s3.amazonaws.com/models.huggingface.co/reformer/reformer-base-cased-pytorch_model.bin",
# "reformer-large-cased": "https://s3.amazonaws.com/models.huggingface.co/reformer/reformer-large-cased-pytorch_model.bin",
PATH_TO_SAVE_WEIGHTS = "/home/patrick/hugging_face/experiments/reformer/intermediate_weights"
def mish(x):
return x * torch.tanh(nn.functional.softplus(x))
ACT2FN = {"gelu": gelu, "relu": torch.nn.functional.relu, "swish": swish, "gelu_new": gelu_new, "mish": mish}
ReformerLayerNorm = torch.nn.LayerNorm
def create_sinusoidal_embeddings(n_pos, dim, out):
position_enc = np.array([[pos / np.power(10000, 2 * (j // 2) / dim) for j in range(dim)] for pos in range(n_pos)])
out[:, 0::2] = torch.FloatTensor(np.sin(position_enc[:, 0::2]))
out[:, 1::2] = torch.FloatTensor(np.cos(position_enc[:, 1::2]))
out.detach_()
out.requires_grad = False
class ReformerEmbeddings(nn.Module):
"""Construct the embeddings from word, position and token_type embeddings.
"""
def __init__(self, config):
super().__init__()
self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size)
self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)
if config.sinusoidal_pos_embds:
create_sinusoidal_embeddings(
n_pos=config.max_position_embeddings, dim=config.hidden_size, out=self.position_embeddings.weight
)
self.dropout = config.hidden_dropout_prob
def forward(self, input_ids=None, position_ids=None, inputs_embeds=None):
if input_ids is not None:
input_shape = input_ids.size()
else:
input_shape = inputs_embeds.size()[:-1]
seq_length = input_shape[1]
device = input_ids.device if input_ids is not None else inputs_embeds.device
if position_ids is None:
position_ids = torch.arange(seq_length, dtype=torch.long, device=device)
position_ids = position_ids.unsqueeze(0).expand(input_shape)
if inputs_embeds is None:
inputs_embeds = self.word_embeddings(input_ids)
position_embeddings = self.position_embeddings(position_ids)
# TODO: In https://github.com/google/trax/blob/d947ebc184b94bc48d7e48cc63156a047f00992e/trax/models/reformer/reformer.py#L775 there are two dropouts right after each other as
# it is implemented here. Check if this is ok.
embeddings = nn.functional.dropout(inputs_embeds, self.dropout, self.training)
embeddings = embeddings + position_embeddings
embeddings = nn.functional.dropout(embeddings, self.dropout, self.training)
return embeddings
class LSHSelfAttention(nn.Module):
def __init__(self, config):
super().__init__()
self.num_attention_heads = config.num_attention_heads
self.hidden_size = config.hidden_size
self.hash_seed = config.seed
self.num_hashes = config.num_hashes
self.num_buckets = config.num_buckets
self.chunk_length = config.chunk_length
self.num_chunks_before = config.num_chunks_before
self.num_chunks_after = config.num_chunks_after
self.output_attentions = config.output_attentions
self.is_decoder = config.is_decoder
self.attention_head_size = int(self.hidden_size / self.num_attention_heads)
self.all_head_size = self.num_attention_heads * self.attention_head_size
if self.hidden_size % self.num_attention_heads != 0:
raise ValueError(
"The hidden size (%d) is not a multiple of the number of attention "
"heads (%d)" % (config.hidden_size, config.num_attention_heads)
)
self.query_key = nn.Linear(self.hidden_size, self.all_head_size, bias=False)
self.value = nn.Linear(self.hidden_size, self.all_head_size, bias=False)
self.dropout = config.attention_probs_dropout_prob
def forward(
self,
hidden_states,
head_mask=None,
):
# get SeqLen and BatchSize
sequence_length = hidden_states.shape[1]
batch_size = hidden_states.shape[0]
mixed_query_key_vectors = self.query_key(hidden_states)
mixed_value_vectors = self.value(hidden_states)
query_key_vectors = self._transpose_for_scores(mixed_query_key_vectors)
value_vectors = self._transpose_for_scores(mixed_value_vectors)
assert query_key_vectors.shape[-1] == self.attention_head_size
assert value_vectors.shape[-1] == self.attention_head_size
# hash query key vectors into buckets
buckets = self._hash_vectors(query_key_vectors)
assert int(buckets.shape[-1]) == self.num_hashes * sequence_length
ticker, undo_ticker = self._get_ticker_and_undo_ticker(sequence_length, buckets)
query_key_vectors = self._gather_by_expansion(query_key_vectors, ticker)
value_vectors = self._gather_by_expansion(value_vectors, ticker)
# q_info = ticker
query_key_vectors = self._split_dim_by(query_key_vectors, -1, self.chunk_length)
value_vectors = self._split_dim_by(value_vectors, -1, self.chunk_length)
ticker = self._split_dim_by(ticker, -1, self.chunk_length)
# Optionally include adjacent chunks.
if self.chunk_length is None:
assert self.num_chunks_before == 0 and self.num_chunks_after == 0
key_vectors = self._len_and_dim_norm(query_key_vectors)
out_vectors, logits, attention_probs = self._attend(query_key_vectors, key_vectors, value_vectors, ticker, undo_ticker, head_mask)
if self.num_hashes > 1:
out_vectors = self._split_dim_by(out_vectors, self.num_hashes, sequence_length)
logits = self._split_dim_by(logits, self.num_hashes, sequence_length).unsqueeze(-1)
probs_vectors = torch.exp(logits - torch.logsumexp(logits, dim=2, keepdim=True))
out_vectors = torch.sum(out_vectors * probs_vectors, dim=2)
assert out_vectors.shape == (batch_size, self.num_attention_heads, sequence_length, self.attention_head_size)
out_vectors = self._transpose_for_output(out_vectors)
outputs = (out_vectors, attention_probs) if self.output_attentions else (out_vectors,)
return outputs
def _hash_vectors(self, vectors):
batch_size = vectors.shape[0]
# See https://arxiv.org/pdf/1509.02897.pdf
# We sample a different random rotation for each round of hashing to
# decrease the probability of hash misses.
if isinstance(self.num_buckets, int):
assert self.num_buckets % 2 == 0, 'There should be an even number of bucktes, but `self.num_bucktes`: {}'.format(self.num_buckets)
rotation_size = self.num_buckets
num_buckets = self.num_buckets
else:
# Factorize the hash if self.num_buckets is a list or tuple
rotation_size, num_buckets = 0, 1
for num_bucket in self.num_buckets:
assert num_bucket % 2 == 0, 'The number of buckets should be even, but `num_bucket`: {}'.format(num_bucket)
rotation_size += num_bucket
num_buckets *= num_bucket
rotations_shape = (vectors.shape[-1], self.num_hashes, rotation_size // 2)
# TODO: delete later when integration tests are ok
# create a random self.attention_head_size x self.num_hashes x self.num_buckets/2
# random_rotations = torch.randn(rotations_shape, device=vectors.device)
numpy.random.seed(self.hash_seed)
random_rotations = torch.tensor(numpy.random.normal(size=rotations_shape), dtype=torch.float32, device=vectors.device)
# rotated_vectors has dim:
# Output dim: Batch_Size x Num_Attn_Heads x Num_Hashes x Seq_Len x Num_Buckets/2
# TODO: IMPORTANT: At the moment we use the same random rotation over all batches
# and heads -> is that bad? It seems like in original reformer a different random
# rotation is used over heads and batches
rotated_vectors = torch.einsum('bmtd,dhr->bmhtr', vectors, random_rotations)
if isinstance(self.num_buckets, int) or len(self.num_buckets) == 1:
rotated_vectors = torch.cat([rotated_vectors, -rotated_vectors], dim=-1)
buckets = torch.argmax(rotated_vectors, dim=-1)
else:
# Get the buckets for them and combine.
buckets, cur_sum, cur_product = None, 0, 1
for num_bucket in self.num_buckets:
rotated_vectors = rotated_vectors[..., cur_sum:cur_sum + (num_bucket // 2)]
cur_sum += num_bucket // 2
rotated_vectors = torch.cat([rotated_vectors, -rotated_vectors], dim=-1)
if buckets is None:
buckets = torch.argmax(rotated_vectors, dim=-1)
else:
buckets += cur_product * torch.argmax(rotated_vectors, dim=-1)
cur_product *= num_bucket
# buckets is now (Batch_size x Num_Attn_Heads x Num_Hashes x Seq_Len).
# Next we add offsets so that bucket numbers from different hashing rounds don't overlap.
offsets = torch.arange(self.num_hashes, device=vectors.device)
offsets = torch.reshape(offsets * num_buckets, (-1, 1))
# repeat same values for Batch_size and Num_Attn_Heads
offsets = offsets.repeat(batch_size, self.num_attention_heads, 1, 1)
offset_buckets = self._merge_by_middle_dim(buckets + offsets)
return offset_buckets
def _get_ticker_and_undo_ticker(self, sequence_length, buckets):
batch_size = buckets.shape[0]
# TODO: what is ticker? Is ticker something like indices?? Ask authors
ticker = torch.arange(self.num_hashes * sequence_length, device=buckets.device)
ticker = ticker.repeat(batch_size, self.num_attention_heads, 1)
buckets_and_t = sequence_length * buckets + (ticker % sequence_length)
# Hash-based sort
sorted_ticker = torch.argsort(buckets_and_t, dim=-1)
undo_sorted_ticker = torch.argsort(sorted_ticker, dim=-1)
sorted_ticker = (sorted_ticker % sequence_length)
return sorted_ticker, undo_sorted_ticker
def _attend(self, query_vectors, key_vectors, value_vectors, ticker, undo_ticker, head_mask):
key_vectors = self._look_adjacent(key_vectors)
value_vectors = self._look_adjacent(value_vectors)
# get logits and dots
query_key_dots = torch.matmul(query_vectors, key_vectors.transpose(-1, -2))
# TODO(PVP): better naming
query_info = ticker
key_value_info = self._look_adjacent(ticker)
# Causal mask
if self.is_decoder:
# TODO (PVP): This line can be improved in terms of memory. Mask should be inserted and does not have to be created each time layer is called
mask = torch.lt(query_info.unsqueeze(-1), key_value_info.unsqueeze(-2)).long().to(query_info.device)
query_key_dots = query_key_dots - mask * 1e9
# Self mask
# TODO (PVP): This line can be improved in terms of memory. Mask should be inserted and does not have to be created each time layer is called
mask = torch.eq(query_info.unsqueeze(-1), key_value_info.unsqueeze(-2)).long().to(query_info.device)
query_key_dots = query_key_dots - mask * 1e5
# Note: Causal mask probably uses higher mask value (-1e9) than Self mask (-1e5) so that token is able to attend to itself when it has no other valid attention targets.
# Note: Self mask is used because Q and K projection weights are shared.
# From the reformer paper (https://arxiv.org/pdf/2001.04451.pdf):
# " While attention to the future is not allowed, typical implementations of the
# Transformer do allow a position to attend to itself.
# Such behavior is undesirable in a shared-QK formulation because the dot-product
# of a query vector with itself will almost always be greater than the dot product of a
# query vector with a vector at another position. We therefore modify the masking
# to forbid a token from attending to itself, except in situations
# where a token has no other valid attention targets (e.g. the first token in a sequence) "
logits = torch.logsumexp(query_key_dots, dim=-1, keepdim=True)
dots = torch.exp(query_key_dots - logits)
# dropout
dots = nn.functional.dropout(dots, self.dropout, self.training)
# Mask heads if we want to
if head_mask is not None:
dots = dots * head_mask
# attend values
out_vectors = torch.matmul(dots, value_vectors)
# merge chunk length
logits = self._merge_by_middle_dim(logits).squeeze(-1)
out_vectors = self._merge_by_middle_dim(out_vectors)
expanded_undo_sort_indices = undo_ticker.unsqueeze(-1).expand(-1, -1, -1, self.attention_head_size)
out_vectors = torch.gather(out_vectors, 2, expanded_undo_sort_indices)
logits = torch.gather(logits, 2, undo_ticker)
return out_vectors, logits, dots
def _look_adjacent(self, vectors):
""" Used to implement attention between consecutive chunks.
Args:
vectors: array of shape [batch_size, num_attention_heads, n_chunks, chunk_len, ...]
Returns:
array of shape [n_chunks, N * chunk_len, ...], where
N = (1 + n_chunks_before + n_chunks_after).
"""
if self.num_chunks_before == 0 and self.num_chunks_after == 0:
return vectors
slices = []
for i in range(-self.num_chunks_before, self.num_chunks_after + 1):
if i == 0:
slices.append(vectors)
else:
slices.append(torch.cat([vectors[:, :, i:, ...], vectors[:, :, :i, ...]], dim=2))
return torch.cat(slices, dim=3)
def _len_and_dim_norm(self, vectors):
vectors = self._len_norm(vectors)
vectors = vectors / torch.sqrt(torch.tensor(self.attention_head_size, device=vectors.device, dtype=torch.float32))
return vectors
def _len_norm(self, x, epsilon=1e-6):
variance = torch.mean(x**2, -1, keepdim=True)
norm_x = x / torch.sqrt(variance + epsilon)
return norm_x
def _gather_by_expansion(self, vectors, idxs):
expanded_idxs = idxs.unsqueeze(-1).expand(-1, -1, -1, self.attention_head_size)
vectors = vectors.repeat(1, 1, self.num_hashes, 1)
return torch.gather(vectors, 2, expanded_idxs)
def _transpose_for_scores(self, x):
new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
x = x.view(*new_x_shape)
return x.transpose(2, 1)
def _transpose_for_output(self, x):
x = x.permute(0, 2, 1, 3)
return torch.reshape(x, (x.size()[0], -1, self.num_attention_heads * self.attention_head_size))
def _split_dim_by(self, vectors, dim_factor_1, dim_factor_2):
batch_size = vectors.shape[0]
split_dim_shape = (batch_size, self.num_attention_heads, dim_factor_1, dim_factor_2)
if len(vectors.shape) == 4:
return torch.reshape(vectors, split_dim_shape + (self.attention_head_size,))
elif len(vectors.shape) == 3:
return torch.reshape(vectors, split_dim_shape)
else:
raise ValueError("Input vector rank should be one of [3, 4], but is: {}".format(len(vectors.shape)))
def _merge_by_middle_dim(self, vectors):
batch_size = vectors.shape[0]
new_dim_shape = (batch_size, self.num_attention_heads, -1)
if len(vectors.shape) == 5:
return torch.reshape(vectors, new_dim_shape + (vectors.shape[-1],))
elif len(vectors.shape) == 4:
return torch.reshape(vectors, new_dim_shape)
else:
raise ValueError("Input vector rank should be one of [4, 5], but is: {}".format(len(vectors.shape)))
class ReformerSelfOutput(nn.Module):
def __init__(self, config):
super().__init__()
self.dense = nn.Linear(config.hidden_size, config.hidden_size, bias=False)
self.dropout = config.hidden_dropout_prob
def forward(self, hidden_states, input_tensor):
hidden_states = self.dense(hidden_states)
hidden_states = nn.functional.dropout(hidden_states, self.dropout, self.training)
# residual connection
output = (hidden_states + input_tensor)
return output
class ReformerAttention(nn.Module):
def __init__(self, config):
super().__init__()
self.layer_norm = ReformerLayerNorm(config.hidden_size, eps=config.layer_norm_eps)
self.self_attention = LSHSelfAttention(config)
self.output = ReformerSelfOutput(config)
def forward(
self,
hidden_states,
prev_attention_output,
head_mask=None,
):
norm_hidden_states = self.layer_norm(hidden_states)
self_attention_outputs = self.self_attention(
norm_hidden_states, head_mask
)
# Implementation of RevNet (see Fig. 6 in https://towardsdatascience.com/illustrating-the-reformer-393575ac6ba0)
# X_1 = prev_attention_output; X_2 = hidden_states
attention_output = self.output(self_attention_outputs[0], prev_attention_output)
outputs = (attention_output,) + self_attention_outputs[1:] # add attentions if we output them
return outputs
class ReformerFeedForwardDense(nn.Module):
def __init__(self, config):
super().__init__()
self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
self.dropout = config.hidden_dropout_prob
if isinstance(config.hidden_act, str):
self.act_fn = ACT2FN[config.hidden_act]
else:
self.act_fn = config.hidden_act
def forward(self, hidden_states):
hidden_states = self.dense(hidden_states)
hidden_states = nn.functional.dropout(hidden_states, self.dropout, self.training)
hidden_states = self.act_fn(hidden_states)
return hidden_states
class ReformerFeedForwardOutput(nn.Module):
def __init__(self, config):
super().__init__()
self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
self.dropout = config.hidden_dropout_prob
def forward(self, hidden_states, input_tensor):
hidden_states = self.dense(hidden_states)
hidden_states = nn.functional.dropout(hidden_states, self.dropout, self.training)
output = (hidden_states + input_tensor)
return output
class ReformerFeedForward(nn.Module):
def __init__(self, config):
super().__init__()
self.layer_norm = ReformerLayerNorm(config.hidden_size, eps=config.layer_norm_eps)
self.dense = ReformerFeedForwardDense(config)
self.output = ReformerFeedForwardOutput(config)
def forward(self, attention_output, prev_attention_output):
norm_attention_output = self.layer_norm(attention_output)
dense_output = self.dense(norm_attention_output)
output = self.output(dense_output, prev_attention_output)
return output
class ReformerLayer(nn.Module):
def __init__(self, config):
super().__init__()
self.attention = ReformerAttention(config)
self.feed_forward = ReformerFeedForward(config)
def forward(
self,
prev_attention_output,
hidden_states,
head_mask=None,
):
# Implementation of RevNet (see Fig. 6 in https://towardsdatascience.com/illustrating-the-reformer-393575ac6ba0)
# X_1 = prev_attention_output; X_2 = hidden_states
# Y_1 = attention_output; Y_2 = output
attention_outputs = self.attention(hidden_states, prev_attention_output, head_mask)
attention_output = attention_outputs[0]
output = self.feed_forward(attention_output, hidden_states)
outputs = (attention_output, output) + attention_outputs[1:]
return outputs
class ReformerEncoder(nn.Module):
def __init__(self, config):
super().__init__()
self.output_attentions = config.output_attentions
self.output_hidden_states = config.output_hidden_states
self.layer = nn.ModuleList([ReformerLayer(config) for _ in range(config.num_hidden_layers)])
# Reformer is using Rev Nets, thus last layer outputs are concatenated and
# Layer Norm is done over 2 * hidden_size
self.layer_norm = ReformerLayerNorm(2 * config.hidden_size, eps=config.layer_norm_eps)
self.dropout = config.hidden_dropout_prob
def forward(
self,
hidden_states,
head_mask=None,
):
all_hidden_states = ()
all_attentions = ()
attention_output = hidden_states
for i, layer_module in enumerate(self.layer):
if self.output_hidden_states:
all_hidden_states = all_hidden_states + (attention_output, hidden_states)
layer_outputs = layer_module(attention_output, hidden_states, head_mask[i])
attention_output, hidden_states = layer_outputs[:2]
if self.output_attentions:
all_attentions = all_attentions + (layer_outputs[2],)
# Add last layer
if self.output_hidden_states:
all_hidden_states = all_hidden_states + (attention_output, hidden_states)
trax_2_layer_output_1 = np.load(PATH_TO_SAVE_WEIGHTS + '/encoder_output_1.npy')
trax_2_layer_output_2 = np.load(PATH_TO_SAVE_WEIGHTS + '/encoder_output_2.npy')
import ipdb
ipdb.set_trace()
# Concatenate 2 RevNet outputs
hidden_states = torch.cat([attention_output, hidden_states], dim=-1)
# Apply layer norm to concatenated hidden states
hidden_states = self.layer_norm(hidden_states)
# Apply dropout
hidden_states = nn.functional.dropout(hidden_states, self.dropout, self.training)
outputs = (hidden_states,)
if self.output_hidden_states:
outputs = outputs + (all_hidden_states,)
if self.output_attentions:
outputs = outputs + (all_attentions,)
return outputs # last-layer hidden state, (all hidden states), (all attentions)
class ReformerPreTrainedModel(PreTrainedModel):
""" An abstract class to handle weights initialization and
a simple interface for downloading and loading pretrained models.
"""
config_class = ReformerConfig
pretrained_model_archive_map = REFORMER_PRETRAINED_MODEL_ARCHIVE_MAP
base_model_prefix = "reformer"
def _init_weights(self, module):
""" Initialize the weights """
if isinstance(module, (nn.Linear, nn.Embedding)):
# Slightly different from the TF version which uses truncated_normal for initialization
# cf https://github.com/pytorch/pytorch/pull/5617
module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
elif isinstance(module, ReformerLayerNorm):
module.bias.data.zero_()
module.weight.data.fill_(1.0)
if isinstance(module, nn.Linear) and module.bias is not None:
module.bias.data.zero_()
class ReformerModel(ReformerPreTrainedModel):
"""
The model can behave as an encoder (with only self-attention) as well
as a decoder, in which case a layer of cross-attention is added between
the self-attention layers, following the architecture described in `Attention is all you need`_ by Ashish Vaswani,
Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser and Illia Polosukhin.
To behave as an decoder the model needs to be initialized with the
:obj:`is_decoder` argument of the configuration set to :obj:`True`; an
:obj:`encoder_hidden_states` is expected as an input to the forward pass.
.. _`Attention is all you need`:
https://arxiv.org/abs/1706.03762
"""
def __init__(self, config):
super().__init__(config)
self.config = config
self.embeddings = ReformerEmbeddings(config)
self.encoder = ReformerEncoder(config)
self.init_weights()
def get_input_embeddings(self):
return self.embeddings.word_embeddings
def set_input_embeddings(self, value):
self.embeddings.word_embeddings = value
def _prune_heads(self, heads_to_prune):
""" Prunes heads of the model.
heads_to_prune: dict of {layer_num: list of heads to prune in this layer}
See base class PreTrainedModel
"""
for layer, heads in heads_to_prune.items():
self.encoder.layer[layer].attention.prune_heads(heads)
def forward(
self,
input_ids=None,
attention_mask=None,
position_ids=None,
head_mask=None,
inputs_embeds=None,
):
r"""
Return:
:obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~transformers.BertConfig`) and inputs:
last_hidden_state (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`):
Sequence of hidden-states at the output of the last layer of the model.
hidden_states (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_hidden_states=True``):
Tuple of :obj:`torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer)
of shape :obj:`(batch_size, sequence_length, hidden_size)`.
Hidden-states of the model at the output of each layer plus the initial embedding outputs.
attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
:obj:`(batch_size, num_heads, sequence_length, sequence_length)`.
Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
heads.
Examples::
"""
if input_ids is not None and inputs_embeds is not None:
raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
elif input_ids is not None:
input_shape = input_ids.size() # noqa: F841
elif inputs_embeds is not None:
input_shape = inputs_embeds.size()[:-1] # noqa: F841
else:
raise ValueError("You have to specify either input_ids or inputs_embeds")
device = input_ids.device if input_ids is not None else inputs_embeds.device # noqa: F841
# Prepare head mask if needed
# 1.0 in head_mask indicate we keep the head
# attention_probs has shape bsz x n_heads x N x N
# input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
# and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
if head_mask is not None:
if head_mask.dim() == 1:
head_mask = head_mask.unsqueeze(0).unsqueeze(0).unsqueeze(-1).unsqueeze(-1)
head_mask = head_mask.expand(self.config.num_hidden_layers, -1, -1, -1, -1)
elif head_mask.dim() == 2:
head_mask = (
head_mask.unsqueeze(1).unsqueeze(-1).unsqueeze(-1)
) # We can specify head_mask for each layer
head_mask = head_mask.to(
dtype=next(self.parameters()).dtype
) # switch to fload if need + fp16 compatibility
else:
head_mask = [None] * self.config.num_hidden_layers
embedding_output = self.embeddings(
input_ids=input_ids, position_ids=position_ids, inputs_embeds=inputs_embeds
)
encoder_outputs = self.encoder(
embedding_output,
head_mask=head_mask,
)
sequence_output = encoder_outputs[0]
# add hidden_states and attentions if they are here
outputs = (sequence_output,) + encoder_outputs[1:]
return outputs # sequence_output, pooled_output, (hidden_states), (attentions)
class ReformerWithLMHead(ReformerPreTrainedModel):
def __init__(self, config):
super().__init__(config)
self.reformer = ReformerModel(config)
self.lm_head = ReformerOnlyLMHead(config)
self.init_weights()
def get_output_embeddings(self):
return self.lm_head.predictions.decoder
@add_start_docstrings_to_callable(BERT_INPUTS_DOCSTRING)
def forward(
self,
input_ids=None,
position_ids=None,
head_mask=None,
inputs_embeds=None,
lm_labels=None,
):
outputs = self.reformer(
input_ids,
attention_mask=attention_mask,
token_type_ids=token_type_ids,
position_ids=position_ids,
head_mask=head_mask,
inputs_embeds=inputs_embeds,
encoder_hidden_states=encoder_hidden_states,
encoder_attention_mask=encoder_attention_mask,
)
sequence_output = outputs[0]
prediction_scores = self.cls(sequence_output)
outputs = (prediction_scores,) + outputs[2:] # Add hidden states and attention if they are here
# Although this may seem awkward, BertForMaskedLM supports two scenarios:
# 1. If a tensor that contains the indices of masked labels is provided,
# the cross-entropy is the MLM cross-entropy that measures the likelihood
# of predictions for masked words.
# 2. If `lm_labels` is provided we are in a causal scenario where we
# try to predict the next token for each input in the decoder.
if masked_lm_labels is not None:
loss_fct = CrossEntropyLoss() # -100 index = padding token
masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), masked_lm_labels.view(-1))
outputs = (masked_lm_loss,) + outputs
if lm_labels is not None:
# we are doing next-token prediction; shift prediction scores and input ids by one
prediction_scores = prediction_scores[:, :-1, :].contiguous()
lm_labels = lm_labels[:, 1:].contiguous()
loss_fct = CrossEntropyLoss()
ltr_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), lm_labels.view(-1))
outputs = (ltr_lm_loss,) + outputs
return outputs # (masked_lm_loss), (ltr_lm_loss), prediction_scores, (hidden_states), (attentions)
tests/test_modeling_reformer.py
+141 -14
View File
@@ -24,9 +24,10 @@ from trax.layers.research.efficient_attention_v2 import (
LSHSelfAttention as TraxLSHSelfAttention,
)
from trax.models.reformer.reformer import DecoderBlock as TraxLSHAttentionBlock
from trax.models.reformer.reformer import ReformerLM as TraxReformer
from trax import layers as tl
from transformers import ReformerAttention, ReformerLayer, ReformerConfig
from transformers import ReformerAttention, ReformerLayer, ReformerConfig, ReformerModelWithLMHead
from transformers import is_torch_available # noqa: F401
@@ -64,11 +65,11 @@ class TraxUtils(object):
def convert_to_jax_array(self, np_array):
return jax.numpy.asarray(np_array)
def get_input_signature(self, shape=None):
def get_input_signature(self, shape=None, dtype=trax_math.numpy.float32):
with trax_math.use_backend("jax"):
if shape is None:
shape = self._shape
input_signature = trax_ShapeDtype(shape)
input_signature = trax_ShapeDtype(shape, dtype)
return input_signature
def get_layer(
@@ -163,9 +164,9 @@ class TraxUtils(object):
hash_seed=config.seed,
path_to_save_weights=path_to_save_weights
)
layer = tl.Serial(tl.ReversibleSerial([list_of_layers]))
block = tl.Serial(tl.ReversibleSerial([list_of_layers]))
return layer
return block
def forward_block(
self,
@@ -193,6 +194,79 @@ class TraxUtils(object):
return output, weights, state
def get_model(
self,
config,
use_reference_code=True,
share_qk=True,
ff_use_sru=0,
axial_pos_shape=(),
d_axial_pos_embs=None,
mode="eval",
path_to_save_weights=PATH_TO_SAVE_WEIGHTS,
):
with trax_math.use_backend("jax"):
with jax.disable_jit():
hidden_size_per_head = config.hidden_size // config.num_attention_heads
model = TraxReformer(
vocab_size=config.vocab_size,
d_model=config.d_model,
d_ff=config.d_ff,
d_attention_key=hidden_size_per_head,
d_attention_value=hidden_size_per_head,
n_layers=config.num_hidden_layers,
n_heads=config.num_attention_heads,
dropout=config.hidden_dropout_prob,
max_len=config.max_position_embeddings,
n_chunks=config.num_chunks,
n_attention_chunks=config.num_attention_chunks,
attention_type=tl.LSHSelfAttention,
share_qk=share_qk,
axial_pos_shape=axial_pos_shape,
d_axial_pos_embs=d_axial_pos_embs,
ff_activation=tl.Gelu,
ff_use_sru=ff_use_sru,
ff_chunk_size=config.ff_chunk_size,
mode=mode,
causal=config.is_decoder,
chunk_len=config.chunk_length,
n_chunks_before=config.num_chunks_before,
n_chunks_after=config.num_chunks_after,
n_hashes=config.num_hashes,
n_buckets=config.num_buckets,
use_reference_code=use_reference_code,
hash_seed=config.seed,
path_to_save_weights=path_to_save_weights
)
return model
def forward_model(
self,
np_input_data,
model,
input_signature=None,
random_number_generator=None,
):
with trax_math.use_backend("jax"):
input_data = self.convert_to_jax_array(np_input_data)
input_data = (input_data,) * 2
if input_signature is None:
input_signature = self.get_input_signature(dtype=trax_math.numpy.int32)
input_signature = (input_signature, input_signature)
weights, state = model.init(input_signature)
if random_number_generator is None:
random_number_generator = model.new_rngs(1)[0]
output = model(
input_data, weights=weights, state=state, rng=random_number_generator
)
return output, weights, state
@require_torch
class ReformerIntegrationTests(unittest.TestCase):
@@ -215,18 +289,17 @@ class ReformerIntegrationTests(unittest.TestCase):
self._set_param(torch_layer.self_attention.value, torch.tensor(np_value).transpose(1, 2).contiguous().view(-1, hidden_size))
self._set_param(torch_layer.output.dense, torch.tensor(np_dense).view(-1, hidden_size).contiguous().transpose(0, 1))
def _set_block_weights_in_torch(self, weights, torch_layer, hidden_size):
weights = weights[0]
def _set_block_weights_in_torch(self, weights, torch_block, hidden_size):
# layernorm 1
layer_norm_1 = weights[0][0][0]
layer_norm_1_weight = np.asarray(layer_norm_1[0])
layer_norm_1_bias = np.asarray(layer_norm_1[1])
self._set_param(torch_layer.attention.layer_norm, torch.tensor(layer_norm_1_weight), torch.tensor(layer_norm_1_bias))
self._set_param(torch_block.attention.layer_norm, torch.tensor(layer_norm_1_weight), torch.tensor(layer_norm_1_bias))
# lsh weights + output
lsh_weights = weights[0][1]
self._set_layer_weights_in_torch(lsh_weights, torch_layer.attention, hidden_size)
self._set_layer_weights_in_torch(lsh_weights, torch_block.attention, hidden_size)
# intermediate weighs
intermediate_weights = weights[2][0][2][2]
@@ -234,17 +307,51 @@ class ReformerIntegrationTests(unittest.TestCase):
# layernorm 2
layer_norm_2_weight = np.asarray(intermediate_weights[0][0])
layer_norm_2_bias = np.asarray(intermediate_weights[0][1])
self._set_param(torch_layer.feed_forward.layer_norm, torch.tensor(layer_norm_2_weight), torch.tensor(layer_norm_2_bias))
self._set_param(torch_block.feed_forward.layer_norm, torch.tensor(layer_norm_2_weight), torch.tensor(layer_norm_2_bias))
# intermediate dense
inter_dense_weight = np.asarray(intermediate_weights[1][0])
inter_dense_bias = np.asarray(intermediate_weights[1][1])
self._set_param(torch_layer.feed_forward.dense.dense, torch.tensor(inter_dense_weight).transpose(0, 1).contiguous(), torch.tensor(inter_dense_bias))
self._set_param(torch_block.feed_forward.dense.dense, torch.tensor(inter_dense_weight).transpose(0, 1).contiguous(), torch.tensor(inter_dense_bias))
# intermediate out
out_dense_weight = np.asarray(intermediate_weights[4][0])
out_dense_bias = np.asarray(intermediate_weights[4][1])
self._set_param(torch_layer.feed_forward.output.dense, torch.tensor(out_dense_weight).transpose(0, 1).contiguous(), torch.tensor(out_dense_bias))
self._set_param(torch_block.feed_forward.output.dense, torch.tensor(out_dense_weight).transpose(0, 1).contiguous(), torch.tensor(out_dense_bias))
def _set_model_weights_in_torch(self, weights, torch_model, hidden_size):
# reformer model
torch_model_reformer = torch_model.reformer
# word embeds
word_embeddings = np.asarray(weights[1])
self._set_param(torch_model_reformer.embeddings.word_embeddings, torch.tensor(word_embeddings))
# pos embeds
pos_embeddings = np.asarray(weights[3])
self._set_param(torch_model_reformer.embeddings.position_embeddings, torch.tensor(pos_embeddings[0]))
trax_layer_weights = weights[5]
assert len(torch_model_reformer.encoder.layer) * 4 + 1 == len(trax_layer_weights), "HF and trax model do not have the same number of layers"
for layer_idx, layer in enumerate(torch_model_reformer.encoder.layer):
block_weights = trax_layer_weights[4 * layer_idx: 4 * (layer_idx + 1)]
self._set_block_weights_in_torch(block_weights, layer, hidden_size)
# output weights
out_weights = weights[6]
# output layer norm
layer_norm_out_weight = np.asarray(out_weights[0][0])
layer_norm_out_bias = np.asarray(out_weights[0][1])
self._set_param(torch_model_reformer.encoder.layer_norm, torch.tensor(layer_norm_out_weight), torch.tensor(layer_norm_out_bias))
# output embeddings
output_embed_weights = np.asarray(out_weights[2][0])
output_embed_bias = np.asarray(out_weights[2][1])
self._set_param(torch_model.lm_head.decoder, torch.tensor(output_embed_weights).transpose(0, 1).contiguous(), torch.tensor(output_embed_bias))
pass
def test_lsh_layer(self):
# Remove residual connection in ReformerSelfOutput to test this layer only
@@ -282,7 +389,7 @@ class ReformerIntegrationTests(unittest.TestCase):
hf_input = torch.tensor(np_input, dtype=torch.float)
hf_block = ReformerLayer(config)
self._set_block_weights_in_torch(trax_weights, hf_block, config.hidden_size)
self._set_block_weights_in_torch(trax_weights[0], hf_block, config.hidden_size)
hf_block.eval()
hf_output_1, hf_output_2 = hf_block(hf_input, hf_input)[:2]
@@ -290,4 +397,24 @@ class ReformerIntegrationTests(unittest.TestCase):
self.assertTrue(torch.allclose(hf_output_1, trax_torch_output_1, atol=1e-3))
self.assertTrue(torch.allclose(hf_output_2, trax_torch_output_2, atol=1e-3))
pass
def test_reformer_lm_model(self):
config = ReformerConfig()
shape = (2, 7) # Batch x SeqLen x ModelDimPerHead
np_input = np.random.randint(0, config.vocab_size, size=shape)
np_zeros = np.zeros((shape[0], 1), dtype=np.int)
trax_utils = TraxUtils(shape)
trax_model = trax_utils.get_model(config)
trax_output, trax_weights, trax_state = trax_utils.forward_model(np_input, model=trax_model)
trax_torch_output = torch.tensor(np.asarray(trax_output[0]))
hf_input = torch.cat([torch.tensor(np_zeros), torch.tensor(np_input[:, :-1])], dim=-1)
hf_model = ReformerModelWithLMHead(config)
self._set_model_weights_in_torch(trax_weights, hf_model, config.hidden_size)
hf_model.eval()
hf_output = hf_model(hf_input)
log_softmax_output = torch.nn.functional.log_softmax(hf_output[0], dim=-1)
self.assertTrue(torch.allclose(log_softmax_output, trax_torch_output, atol=1e-3))