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

almost finish BERT

Browse Source
This commit is contained in:
Patrick von Platen
2020-12-11 13:35:32 +00:00
parent 73797f1916
commit caf25047af
6 changed files with 222 additions and 114 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/transformers/__init__.py
+1 -1
View File
@@ -937,7 +937,7 @@ else:
if is_flax_available():
from .models.auto import FLAX_MODEL_MAPPING, FlaxAutoModel
from .models.bert import FlaxBertForMaskedLM, FlaxBertModel
from .models.bert import FlaxBertForMaskedLM, FlaxBertForMaskedLMModule, FlaxBertModel, FlaxBertModule
from .models.roberta import FlaxRobertaModel
else:
# Import the same objects as dummies to get them in the namespace.
src/transformers/modeling_flax_utils.py
+14 -6
View File
@@ -76,8 +76,8 @@ class FlaxPreTrainedModel(ABC):
if config is None:
raise ValueError("config cannot be None")
if params is None:
raise ValueError("state cannot be None")
if module is None:
raise ValueError("module cannot be None")
# Those are private to be exposed as typed property on derived classes.
self._config = config
@@ -194,11 +194,9 @@ class FlaxPreTrainedModel(ABC):
try:
if from_pt:
import torch
state = torch.load(state_f)
state = {k: v.numpy() for k, v in state.items()}
state = cls.convert_from_pytorch(state, config)
state = unflatten_dict({tuple(k.split(".")[1:]): v for k, v in state.items()})
state = convert_state_dict_from_pt(cls, state, config)
else:
state = from_bytes(cls, state_f.read())
except UnpicklingError:
@@ -230,3 +228,13 @@ class FlaxPreTrainedModel(ABC):
with open(os.path.join(save_directory, FLAX_WEIGHTS_NAME), "wb") as f:
model_bytes = to_bytes(self.params)
f.write(model_bytes)
def convert_state_dict_from_pt(model_class: ABC, state: Dict, config: PretrainedConfig):
"""
Converts a PyTorch parameter state dict to an equivalent Flax parameter state dict
"""
state = {k: v.numpy() for k, v in state.items()}
state = model_class.convert_from_pytorch(state, config)
state = unflatten_dict({tuple(k.split(".")): v for k, v in state.items()})
return state
src/transformers/models/bert/__init__.py
+1 -1
View File
@@ -59,4 +59,4 @@ if is_tf_available():
)
if is_flax_available():
from .modeling_flax_bert import FlaxBertForMaskedLM, FlaxBertModel
from .modeling_flax_bert import FlaxBertForMaskedLM, FlaxBertForMaskedLMModule, FlaxBertModel, FlaxBertModule
src/transformers/models/bert/modeling_flax_bert.py
+155 -85
View File
@@ -398,6 +398,87 @@ class FlaxBertPreTrainedModel(FlaxPreTrainedModel):
config_class = BertConfig
base_model_prefix = "bert"
def _check_inputs(self, input_ids, attention_mask, token_type_ids, position_ids):
if token_type_ids is None:
token_type_ids = jnp.ones_like(input_ids)
if position_ids is None:
position_ids = jnp.arange(jnp.atleast_2d(input_ids).shape[-1])
if attention_mask is None:
attention_mask = jnp.ones_like(input_ids)
return input_ids, attention_mask, token_type_ids, position_ids
def init(self, rng: jax.random.PRNGKey, input_shape: Tuple):
input_ids, attention_mask, token_type_ids, position_ids = self._check_inputs(
jnp.zeros(input_shape, dtype="i4"), None, None, None
)
params_rng, dropout_rng = jax.random.split(rng)
rngs = {"params": params_rng, "dropout": dropout_rng}
self.params = self.module.init(rngs, input_ids, attention_mask, token_type_ids, position_ids)["params"]
@add_start_docstrings(
"The bare Bert Model transformer outputting raw hidden-states without any specific head on top.",
BERT_START_DOCSTRING,
)
class FlaxBertModel(FlaxBertPreTrainedModel):
"""
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 <https://arxiv.org/abs/1706.03762>`__ by Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit,
Llion Jones, Aidan N. Gomez, Lukasz Kaiser and Illia Polosukhin.
"""
def __init__(self, config: BertConfig, state: dict = None, seed: int = 0, dtype: jnp.dtype = jnp.float32):
module = FlaxBertModule(
vocab_size=config.vocab_size,
hidden_size=config.hidden_size,
type_vocab_size=config.type_vocab_size,
max_length=config.max_position_embeddings,
num_encoder_layers=config.num_hidden_layers,
num_heads=config.num_attention_heads,
head_size=config.hidden_size,
intermediate_size=config.intermediate_size,
dropout_rate=config.hidden_dropout_prob,
dtype=dtype,
)
super().__init__(config, module, state, seed)
@add_start_docstrings_to_model_forward(BERT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
def __call__(
self,
input_ids,
attention_mask=None,
token_type_ids=None,
position_ids=None,
params: dict = None,
dropout_rng: PRNGKey = None,
train: bool = False,
):
input_ids, attention_mask, token_type_ids, position_ids = self._check_inputs(
input_ids, attention_mask, token_type_ids, position_ids
)
# Handle any PRNG if needed
rngs = {}
if dropout_rng is not None:
rngs["dropout"] = dropout_rng
return self.module.apply(
{"params": params or self.params},
jnp.array(input_ids, dtype="i4"),
jnp.array(attention_mask, dtype="i4"),
jnp.array(token_type_ids, dtype="i4"),
jnp.array(position_ids, dtype="i4"),
not train,
rngs=rngs,
)
@staticmethod
def convert_from_pytorch(pt_state: Dict, config: BertConfig) -> Dict:
jax_state = dict(pt_state)
@@ -468,87 +549,6 @@ class FlaxBertPreTrainedModel(FlaxPreTrainedModel):
return jax_state
@add_start_docstrings(
"The bare Bert Model transformer outputting raw hidden-states without any specific head on top.",
BERT_START_DOCSTRING,
)
class FlaxBertModel(FlaxBertPreTrainedModel):
"""
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 <https://arxiv.org/abs/1706.03762>`__ by Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit,
Llion Jones, Aidan N. Gomez, Lukasz Kaiser and Illia Polosukhin.
"""
def __init__(self, config: BertConfig, state: dict, seed: int = 0, dtype: jnp.dtype = jnp.float32):
module = FlaxBertModule(
vocab_size=config.vocab_size,
hidden_size=config.hidden_size,
type_vocab_size=config.type_vocab_size,
max_length=config.max_position_embeddings,
num_encoder_layers=config.num_hidden_layers,
num_heads=config.num_attention_heads,
head_size=config.hidden_size,
intermediate_size=config.intermediate_size,
dropout_rate=config.hidden_dropout_prob,
dtype=dtype,
)
super().__init__(config, module, state, seed)
@add_start_docstrings_to_model_forward(BERT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
def __call__(
self,
input_ids,
attention_mask=None,
token_type_ids=None,
position_ids=None,
params: dict = None,
dropout_rng: PRNGKey = None,
train: bool = False,
):
input_ids, attention_mask, token_type_ids, position_ids = self._check_inputs(
input_ids, attention_mask, token_type_ids, position_ids
)
# Handle any PRNG if needed
rngs = {}
if dropout_rng is not None:
rngs["dropout"] = dropout_rng
return self.module.apply(
{"params": params or self.params},
jnp.array(input_ids, dtype="i4"),
jnp.array(attention_mask, dtype="i4"),
jnp.array(token_type_ids, dtype="i4"),
jnp.array(position_ids, dtype="i4"),
not train,
rngs=rngs,
)
def _check_inputs(self, input_ids, attention_mask, token_type_ids, position_ids):
if token_type_ids is None:
token_type_ids = jnp.ones_like(input_ids)
if position_ids is None:
position_ids = jnp.arange(jnp.atleast_2d(input_ids).shape[-1])
if attention_mask is None:
attention_mask = jnp.ones_like(input_ids)
return input_ids, attention_mask, token_type_ids, position_ids
def init(self, rng: jax.random.PRNGKey, input_shape: Tuple):
input_ids, attention_mask, token_type_ids, position_ids = self._check_inputs(
jnp.zeros(input_shape, dtype="i4"), None, None, None
)
params_rng, dropout_rng = jax.random.split(rng)
rngs = {"params": params_rng, "dropout": dropout_rng}
self.params = self.module.init(rngs, input_ids, attention_mask, token_type_ids, position_ids)["params"]
class FlaxBertModule(nn.Module):
vocab_size: int
hidden_size: int
@@ -594,10 +594,8 @@ class FlaxBertModule(nn.Module):
class FlaxBertForMaskedLM(FlaxBertPreTrainedModel):
def __init__(self, config: BertConfig, state: dict, seed: int = 0, dtype: jnp.dtype = jnp.float32, **kwargs):
super().__init__(config, state, seed, dtype)
self._module = FlaxBertForMaskedLMModule(
def __init__(self, config: BertConfig, state: dict = None, seed: int = 0, dtype: jnp.dtype = jnp.float32, **kwargs):
module = FlaxBertForMaskedLMModule(
vocab_size=config.vocab_size,
type_vocab_size=config.type_vocab_size,
hidden_size=config.hidden_size,
@@ -609,6 +607,8 @@ class FlaxBertForMaskedLM(FlaxBertPreTrainedModel):
**kwargs,
)
super().__init__(config, module, state, seed)
def __call__(
self,
input_ids,
@@ -640,6 +640,75 @@ class FlaxBertForMaskedLM(FlaxBertPreTrainedModel):
return logits, pooled
@staticmethod
def convert_from_pytorch(pt_state: Dict, config: BertConfig) -> Dict:
jax_state = dict(pt_state)
# Need to change some parameters name to match Flax names so that we don't have to fork any layer
for key, tensor in pt_state.items():
# Key parts
key_parts = set(key.split("."))
# Every dense layer has "kernel" parameters instead of "weight"
if "dense.weight" in key:
del jax_state[key]
key = key.replace("weight", "kernel")
jax_state[key] = tensor
# SelfAttention needs also to replace "weight" by "kernel"
if {"query", "key", "value"} & key_parts:
# Flax SelfAttention decomposes the heads (num_head, size // num_heads)
if "bias" in key:
jax_state[key] = tensor.reshape((config.num_attention_heads, -1))
elif "weight":
del jax_state[key]
key = key.replace("weight", "kernel")
tensor = tensor.reshape((config.num_attention_heads, -1, config.hidden_size)).transpose((2, 0, 1))
jax_state[key] = tensor
# SelfAttention output is not a separate layer, remove one nesting
if "attention.output.dense" in key:
del jax_state[key]
key = key.replace("attention.output.dense", "attention.self.out")
jax_state[key] = tensor
# SelfAttention output is not a separate layer, remove nesting on layer norm
if "attention.output.LayerNorm" in key:
del jax_state[key]
key = key.replace("attention.output.LayerNorm", "attention.LayerNorm")
jax_state[key] = tensor
# There are some transposed parameters w.r.t their PyTorch counterpart
if "intermediate.dense.kernel" in key or "output.dense.kernel" in key:
jax_state[key] = tensor.T
# Self Attention output projection needs to be transposed
if "out.kernel" in key:
jax_state[key] = tensor.reshape((config.hidden_size, config.num_attention_heads, -1)).transpose(
1, 2, 0
)
# Pooler needs to transpose its kernel
if "pooler.dense.kernel" in key:
jax_state[key] = tensor.T
# Handle LayerNorm conversion
if "LayerNorm" in key:
del jax_state[key]
# Replace LayerNorm by layer_norm
new_key = key.replace("LayerNorm", "layer_norm")
if "weight" in key:
new_key = new_key.replace("weight", "gamma")
elif "bias" in key:
new_key = new_key.replace("bias", "beta")
jax_state[new_key] = tensor
return jax_state
class FlaxBertForMaskedLMModule(nn.Module):
vocab_size: int
@@ -669,6 +738,7 @@ class FlaxBertForMaskedLMModule(nn.Module):
max_length=self.max_length,
dropout_rate=self.dropout_rate,
dtype=self.dtype,
name="bert"
)(input_ids, attention_mask, token_type_ids, position_ids, deterministic=deterministic)
# Compute the prediction scores
tests/test_modeling_flax_bert.py
+9 -3
View File
@@ -15,13 +15,13 @@
import unittest
from transformers import BertConfig, is_flax_available
from transformers.testing_utils import require_flax
from transformers.testing_utils import require_flax, slow
from .test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask
if is_flax_available():
from transformers.models.bert.modeling_flax_bert import FlaxBertModel
from transformers.models.bert.modeling_flax_bert import FlaxBertModel, FlaxBertForMaskedLM
class FlaxBertModelTester(unittest.TestCase):
@@ -105,7 +105,13 @@ class FlaxBertModelTester(unittest.TestCase):
@require_flax
class FlaxBertModelTest(FlaxModelTesterMixin, unittest.TestCase):
all_model_classes = (FlaxBertModel,) if is_flax_available() else ()
all_model_classes = (FlaxBertModel, FlaxBertForMaskedLM) if is_flax_available() else ()
def setUp(self):
self.model_tester = FlaxBertModelTester(self)
@slow
def test_model_from_pretrained(self):
for model_class_name in self.all_model_classes:
model = model_class_name.from_pretrained("bert-base-cased")
self.assertIsNotNone(model)
tests/test_modeling_flax_common.py
+42 -18
View File
@@ -13,6 +13,7 @@
# limitations under the License.
import random
import tempfile
import numpy as np
@@ -26,7 +27,7 @@ if is_flax_available():
import jax
import jax.numpy as jnp
from flax.traverse_util import unflatten_dict
from transformers.modeling_flax_utils import convert_state_dict_from_pt
os.environ["XLA_PYTHON_CLIENT_MEM_FRACTION"] = "0.12" # assumed parallelism: 8
@@ -59,14 +60,6 @@ def random_attention_mask(shape, rng=None):
return attn_mask
def convert_pt_model_to_flax(pt_model, config, flax_model_cls):
state = pt_model.state_dict()
state = {k: v.numpy() for k, v in state.items()}
state = flax_model_cls.convert_from_pytorch(state, config)
state = unflatten_dict({tuple(k.split(".")): v for k, v in state.items()})
return flax_model_cls(config, state, dtype=jnp.float32)
@require_flax
class FlaxModelTesterMixin:
model_tester = None
@@ -86,7 +79,8 @@ class FlaxModelTesterMixin:
pt_model_class = getattr(transformers, pt_model_class_name)
pt_model = pt_model_class(config).eval()
fx_model = convert_pt_model_to_flax(pt_model, config, model_class)
fx_state = convert_state_dict_from_pt(model_class, pt_model.state_dict(), config)
fx_model = model_class(config, fx_state, dtype=jnp.float32)
pt_inputs = {k: torch.tensor(v.tolist()) for k, v in inputs_dict.items()}
@@ -97,19 +91,40 @@ class FlaxModelTesterMixin:
for fx_output, pt_output in zip(fx_outputs, pt_outputs):
self.assert_almost_equals(fx_output, pt_output.numpy(), 5e-3)
@require_torch
with tempfile.TemporaryDirectory() as tmpdirname:
pt_model.save_pretrained(tmpdirname)
fx_model_loaded = model_class.from_pretrained(tmpdirname, from_pt=True)
fx_outputs_loaded = fx_model_loaded(**inputs_dict)
self.assertEqual(len(fx_outputs_loaded), len(pt_outputs), "Output lengths differ between Flax and PyTorch")
for fx_output_loaded, pt_output in zip(fx_outputs_loaded, pt_outputs):
self.assert_almost_equals(fx_output_loaded, pt_output.numpy(), 5e-3)
def test_from_pretrained_save_pretrained(self):
config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
for model_class in self.all_model_classes:
with self.subTest(model_class.__name__):
model = model_class(config)
model.init(model.key, inputs_dict["input_ids"].shape)
outputs = model(**inputs_dict)
with tempfile.TemporaryDirectory() as tmpdirname:
model.save_pretrained(tmpdirname)
model_loaded = model_class.from_pretrained(tmpdirname)
outputs_loaded = model_loaded(**inputs_dict)
for output_loaded, output in zip(outputs_loaded, outputs):
self.assert_almost_equals(output_loaded, output, 5e-3)
def test_jit_compilation(self):
config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
for model_class in self.all_model_classes:
with self.subTest(model_class.__name__):
# TODO later: have some way to initialize easily a Flax model from config, for now I go through PT
pt_model_class_name = model_class.__name__[4:] # Skip the "Flax" at the beginning
pt_model_class = getattr(transformers, pt_model_class_name)
pt_model = pt_model_class(config).eval()
model = convert_pt_model_to_flax(pt_model, config, model_class)
model = model_class(config)
model.init(model.key, inputs_dict["input_ids"].shape)
@jax.jit
def model_jitted(input_ids, attention_mask=None, token_type_ids=None):
@@ -125,3 +140,12 @@ class FlaxModelTesterMixin:
self.assertEqual(len(outputs), len(jitted_outputs))
for jitted_output, output in zip(jitted_outputs, outputs):
self.assertEqual(jitted_output.shape, output.shape)
def test_naming_convention(self):
for model_class in self.all_model_classes:
model_class_name = model_class.__class__.__name__
module_class_name = model_class_name[:-5] + "Module" if model_class_name[-5:] == "Model" else model_class_name + "Module"
module = __import__("transformers", fromlist=[module_class_name])
self.assertIsNotNone(module)