Aluode/PerceptionLabPortable
0
1# coding=utf-82# Copyright 2021 The HuggingFace Inc. team.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8# http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15"""PyTorch - Flax general utilities."""16 17import os18from pickle import UnpicklingError19 20import jax21import jax.numpy as jnp22import numpy as np23from flax.serialization import from_bytes24from flax.traverse_util import flatten_dict, unflatten_dict25 26import transformers27 28from . import is_safetensors_available, is_torch_available29from .utils import check_torch_load_is_safe, logging30 31 32if is_torch_available():33 import torch34 35if is_safetensors_available():36 from safetensors import safe_open37 from safetensors.flax import load_file as safe_load_file38 39 40logger = logging.get_logger(__name__)41 42 43#####################44# PyTorch => Flax #45#####################46 47 48def load_pytorch_checkpoint_in_flax_state_dict(49 flax_model, pytorch_checkpoint_path, is_sharded, allow_missing_keys=False50):51 """Load pytorch checkpoints in a flax model"""52 53 if not is_sharded:54 pt_path = os.path.abspath(pytorch_checkpoint_path)55 logger.info(f"Loading PyTorch weights from {pt_path}")56 57 if pt_path.endswith(".safetensors"):58 pt_state_dict = {}59 with safe_open(pt_path, framework="flax") as f:60 for k in f.keys():61 pt_state_dict[k] = f.get_tensor(k)62 else:63 try:64 import torch # noqa: F40165 except (ImportError, ModuleNotFoundError):66 logger.error(67 "Loading a PyTorch model in Flax, requires both PyTorch and Flax to be installed. Please see"68 " https://pytorch.org/ and https://flax.readthedocs.io/en/latest/index.html#installation for installation"69 " instructions."70 )71 raise72 73 check_torch_load_is_safe()74 pt_state_dict = torch.load(pt_path, map_location="cpu", weights_only=True)75 logger.info(f"PyTorch checkpoint contains {sum(t.numel() for t in pt_state_dict.values()):,} parameters.")76 77 flax_state_dict = convert_pytorch_state_dict_to_flax(pt_state_dict, flax_model)78 else:79 # model is sharded and pytorch_checkpoint_path already contains the list of .pt shard files80 flax_state_dict = convert_pytorch_sharded_state_dict_to_flax(pytorch_checkpoint_path, flax_model)81 return flax_state_dict82 83 84def rename_key_and_reshape_tensor(85 pt_tuple_key: tuple[str],86 pt_tensor: np.ndarray,87 random_flax_state_dict: dict[str, jnp.ndarray],88 model_prefix: str,89) -> tuple[tuple[str], np.ndarray]:90 """Rename PT weight names to corresponding Flax weight names and reshape tensor if necessary"""91 92 def is_key_or_prefix_key_in_dict(key: tuple[str]) -> bool:93 """Checks if `key` of `(prefix,) + key` is in random_flax_state_dict"""94 return len(set(random_flax_state_dict) & {key, (model_prefix,) + key}) > 095 96 # layer norm97 renamed_pt_tuple_key = pt_tuple_key[:-1] + ("scale",)98 if pt_tuple_key[-1] in ["weight", "gamma"] and is_key_or_prefix_key_in_dict(renamed_pt_tuple_key):99 return renamed_pt_tuple_key, pt_tensor100 101 # batch norm layer mean102 renamed_pt_tuple_key = pt_tuple_key[:-1] + ("mean",)103 if pt_tuple_key[-1] == "running_mean" and not is_key_or_prefix_key_in_dict(pt_tuple_key):104 return renamed_pt_tuple_key, pt_tensor105 106 # batch norm layer var107 renamed_pt_tuple_key = pt_tuple_key[:-1] + ("var",)108 if pt_tuple_key[-1] == "running_var" and not is_key_or_prefix_key_in_dict(pt_tuple_key):109 return renamed_pt_tuple_key, pt_tensor110 111 # embedding112 renamed_pt_tuple_key = pt_tuple_key[:-1] + ("embedding",)113 if pt_tuple_key[-1] == "weight" and is_key_or_prefix_key_in_dict(renamed_pt_tuple_key):114 return renamed_pt_tuple_key, pt_tensor115 116 # conv layer117 renamed_pt_tuple_key = pt_tuple_key[:-1] + ("kernel",)118 if pt_tuple_key[-1] == "weight" and pt_tensor.ndim == 4 and not is_key_or_prefix_key_in_dict(pt_tuple_key):119 pt_tensor = pt_tensor.transpose(2, 3, 1, 0)120 return renamed_pt_tuple_key, pt_tensor121 122 # linear layer123 renamed_pt_tuple_key = pt_tuple_key[:-1] + ("kernel",)124 if pt_tuple_key[-1] == "weight" and not is_key_or_prefix_key_in_dict(pt_tuple_key):125 pt_tensor = pt_tensor.T126 return renamed_pt_tuple_key, pt_tensor127 128 # old PyTorch layer norm weight129 renamed_pt_tuple_key = pt_tuple_key[:-1] + ("weight",)130 if pt_tuple_key[-1] == "gamma":131 return renamed_pt_tuple_key, pt_tensor132 133 # old PyTorch layer norm bias134 renamed_pt_tuple_key = pt_tuple_key[:-1] + ("bias",)135 if pt_tuple_key[-1] == "beta":136 return renamed_pt_tuple_key, pt_tensor137 138 # New `weight_norm` from https://github.com/huggingface/transformers/pull/24030139 name = None140 if pt_tuple_key[-3::2] == ("parametrizations", "original0"):141 name = pt_tuple_key[-2] + "_g"142 elif pt_tuple_key[-3::2] == ("parametrizations", "original1"):143 name = pt_tuple_key[-2] + "_v"144 if name is not None:145 renamed_pt_tuple_key = pt_tuple_key[:-3] + (name,)146 return renamed_pt_tuple_key, pt_tensor147 148 return pt_tuple_key, pt_tensor149 150 151def convert_pytorch_state_dict_to_flax(pt_state_dict, flax_model):152 # convert pytorch tensor to numpy153 from_bin = is_torch_available() and isinstance(next(iter(pt_state_dict.values())), torch.Tensor)154 bfloat16 = torch.bfloat16 if from_bin else "bfloat16"155 156 weight_dtypes = {k: v.dtype for k, v in pt_state_dict.items()}157 158 if from_bin:159 for k, v in pt_state_dict.items():160 # numpy currently does not support bfloat16, need to go over float32 in this case to not lose precision161 if v.dtype == bfloat16:162 v = v.float()163 pt_state_dict[k] = v.cpu().numpy()164 165 model_prefix = flax_model.base_model_prefix166 167 # use params dict if the model contains batch norm layers168 if "params" in flax_model.params:169 flax_model_params = flax_model.params["params"]170 else:171 flax_model_params = flax_model.params172 random_flax_state_dict = flatten_dict(flax_model_params)173 174 # add batch_stats keys,values to dict175 if "batch_stats" in flax_model.params:176 flax_batch_stats = flatten_dict(flax_model.params["batch_stats"])177 random_flax_state_dict.update(flax_batch_stats)178 179 flax_state_dict = {}180 181 load_model_with_head_into_base_model = (model_prefix not in flax_model_params) and (182 model_prefix in {k.split(".")[0] for k in pt_state_dict}183 )184 load_base_model_into_model_with_head = (model_prefix in flax_model_params) and (185 model_prefix not in {k.split(".")[0] for k in pt_state_dict}186 )187 188 # Need to change some parameters name to match Flax names189 for pt_key, pt_tensor in pt_state_dict.items():190 pt_tuple_key = tuple(pt_key.split("."))191 is_bfloat_16 = weight_dtypes[pt_key] == bfloat16192 193 # remove base model prefix if necessary194 has_base_model_prefix = pt_tuple_key[0] == model_prefix195 if load_model_with_head_into_base_model and has_base_model_prefix:196 pt_tuple_key = pt_tuple_key[1:]197 198 # Correctly rename weight parameters199 flax_key, flax_tensor = rename_key_and_reshape_tensor(200 pt_tuple_key, pt_tensor, random_flax_state_dict, model_prefix201 )202 203 # add model prefix if necessary204 require_base_model_prefix = (model_prefix,) + flax_key in random_flax_state_dict205 if load_base_model_into_model_with_head and require_base_model_prefix:206 flax_key = (model_prefix,) + flax_key207 208 if flax_key in random_flax_state_dict:209 if flax_tensor.shape != random_flax_state_dict[flax_key].shape:210 raise ValueError(211 f"PyTorch checkpoint seems to be incorrect. Weight {pt_key} was expected to be of shape "212 f"{random_flax_state_dict[flax_key].shape}, but is {flax_tensor.shape}."213 )214 215 # add batch stats if the model contains batchnorm layers216 if "batch_stats" in flax_model.params:217 if "mean" in flax_key[-1] or "var" in flax_key[-1]:218 flax_state_dict[("batch_stats",) + flax_key] = jnp.asarray(flax_tensor)219 continue220 # remove num_batches_tracked key221 if "num_batches_tracked" in flax_key[-1]:222 flax_state_dict.pop(flax_key, None)223 continue224 225 # also add unexpected weight so that warning is thrown226 flax_state_dict[("params",) + flax_key] = (227 jnp.asarray(flax_tensor) if not is_bfloat_16 else jnp.asarray(flax_tensor, dtype=jnp.bfloat16)228 )229 else:230 # also add unexpected weight so that warning is thrown231 flax_state_dict[flax_key] = (232 jnp.asarray(flax_tensor) if not is_bfloat_16 else jnp.asarray(flax_tensor, dtype=jnp.bfloat16)233 )234 235 return unflatten_dict(flax_state_dict)236 237 238############################239# Sharded Pytorch => Flax #240############################241 242 243def convert_pytorch_sharded_state_dict_to_flax(shard_filenames, flax_model):244 import torch245 246 # Load the index247 flax_state_dict = {}248 for shard_file in shard_filenames:249 # load using msgpack utils250 check_torch_load_is_safe()251 pt_state_dict = torch.load(shard_file, weights_only=True)252 weight_dtypes = {k: v.dtype for k, v in pt_state_dict.items()}253 pt_state_dict = {254 k: v.numpy() if v.dtype != torch.bfloat16 else v.float().numpy() for k, v in pt_state_dict.items()255 }256 257 model_prefix = flax_model.base_model_prefix258 259 # use params dict if the model contains batch norm layers and then add batch_stats keys,values to dict260 if "batch_stats" in flax_model.params:261 flax_model_params = flax_model.params["params"]262 263 random_flax_state_dict = flatten_dict(flax_model_params)264 random_flax_state_dict.update(flatten_dict(flax_model.params["batch_stats"]))265 else:266 flax_model_params = flax_model.params267 random_flax_state_dict = flatten_dict(flax_model_params)268 269 load_model_with_head_into_base_model = (model_prefix not in flax_model_params) and (270 model_prefix in {k.split(".")[0] for k in pt_state_dict}271 )272 load_base_model_into_model_with_head = (model_prefix in flax_model_params) and (273 model_prefix not in {k.split(".")[0] for k in pt_state_dict}274 )275 # Need to change some parameters name to match Flax names276 for pt_key, pt_tensor in pt_state_dict.items():277 pt_tuple_key = tuple(pt_key.split("."))278 is_bfloat_16 = weight_dtypes[pt_key] == torch.bfloat16279 280 # remove base model prefix if necessary281 has_base_model_prefix = pt_tuple_key[0] == model_prefix282 if load_model_with_head_into_base_model and has_base_model_prefix:283 pt_tuple_key = pt_tuple_key[1:]284 285 # Correctly rename weight parameters286 flax_key, flax_tensor = rename_key_and_reshape_tensor(287 pt_tuple_key, pt_tensor, random_flax_state_dict, model_prefix288 )289 # add model prefix if necessary290 require_base_model_prefix = (model_prefix,) + flax_key in random_flax_state_dict291 if load_base_model_into_model_with_head and require_base_model_prefix:292 flax_key = (model_prefix,) + flax_key293 294 if flax_key in random_flax_state_dict:295 if flax_tensor.shape != random_flax_state_dict[flax_key].shape:296 raise ValueError(297 f"PyTorch checkpoint seems to be incorrect. Weight {pt_key} was expected to be of shape "298 f"{random_flax_state_dict[flax_key].shape}, but is {flax_tensor.shape}."299 )300 301 # add batch stats if the model contains batchnorm layers302 if "batch_stats" in flax_model.params:303 if "mean" in flax_key[-1]:304 flax_state_dict[("batch_stats",) + flax_key] = jnp.asarray(flax_tensor)305 continue306 if "var" in flax_key[-1]:307 flax_state_dict[("batch_stats",) + flax_key] = jnp.asarray(flax_tensor)308 continue309 # remove num_batches_tracked key310 if "num_batches_tracked" in flax_key[-1]:311 flax_state_dict.pop(flax_key, None)312 continue313 314 # also add unexpected weight so that warning is thrown315 flax_state_dict[("params",) + flax_key] = (316 jnp.asarray(flax_tensor) if not is_bfloat_16 else jnp.asarray(flax_tensor, dtype=jnp.bfloat16)317 )318 319 else:320 # also add unexpected weight so that warning is thrown321 flax_state_dict[flax_key] = (322 jnp.asarray(flax_tensor) if not is_bfloat_16 else jnp.asarray(flax_tensor, dtype=jnp.bfloat16)323 )324 return unflatten_dict(flax_state_dict)325 326 327#####################328# Flax => PyTorch #329#####################330 331 332def load_flax_checkpoint_in_pytorch_model(model, flax_checkpoint_path):333 """Load flax checkpoints in a PyTorch model"""334 flax_checkpoint_path = os.path.abspath(flax_checkpoint_path)335 logger.info(f"Loading Flax weights from {flax_checkpoint_path}")336 337 # import correct flax class338 flax_cls = getattr(transformers, "Flax" + model.__class__.__name__)339 340 # load flax weight dict341 if flax_checkpoint_path.endswith(".safetensors"):342 flax_state_dict = safe_load_file(flax_checkpoint_path)343 flax_state_dict = unflatten_dict(flax_state_dict, sep=".")344 else:345 with open(flax_checkpoint_path, "rb") as state_f:346 try:347 flax_state_dict = from_bytes(flax_cls, state_f.read())348 except UnpicklingError:349 raise OSError(f"Unable to convert {flax_checkpoint_path} to Flax deserializable object. ")350 351 return load_flax_weights_in_pytorch_model(model, flax_state_dict)352 353 354def load_flax_weights_in_pytorch_model(pt_model, flax_state):355 """Load flax checkpoints in a PyTorch model"""356 357 try:358 import torch # noqa: F401359 except (ImportError, ModuleNotFoundError):360 logger.error(361 "Loading a Flax weights in PyTorch, requires both PyTorch and Flax to be installed. Please see"362 " https://pytorch.org/ and https://flax.readthedocs.io/en/latest/index.html#installation for installation"363 " instructions."364 )365 raise366 367 # check if we have bf16 weights368 is_type_bf16 = flatten_dict(jax.tree_util.tree_map(lambda x: x.dtype == jnp.bfloat16, flax_state)).values()369 if any(is_type_bf16):370 # convert all weights to fp32 if the are bf16 since torch.from_numpy can-not handle bf16371 # and bf16 is not fully supported in PT yet.372 logger.warning(373 "Found ``bfloat16`` weights in Flax model. Casting all ``bfloat16`` weights to ``float32`` "374 "before loading those in PyTorch model."375 )376 flax_state = jax.tree_util.tree_map(377 lambda params: params.astype(np.float32) if params.dtype == jnp.bfloat16 else params, flax_state378 )379 380 flax_state_dict = flatten_dict(flax_state)381 pt_model_dict = pt_model.state_dict()382 383 load_model_with_head_into_base_model = (pt_model.base_model_prefix in flax_state) and (384 pt_model.base_model_prefix not in {k.split(".")[0] for k in pt_model_dict}385 )386 load_base_model_into_model_with_head = (pt_model.base_model_prefix not in flax_state) and (387 pt_model.base_model_prefix in {k.split(".")[0] for k in pt_model_dict}388 )389 390 # keep track of unexpected & missing keys391 unexpected_keys = []392 missing_keys = set(pt_model_dict.keys())393 394 for flax_key_tuple, flax_tensor in flax_state_dict.items():395 has_base_model_prefix = flax_key_tuple[0] == pt_model.base_model_prefix396 require_base_model_prefix = ".".join((pt_model.base_model_prefix,) + flax_key_tuple) in pt_model_dict397 398 # adapt flax_key to prepare for loading from/to base model only399 if load_model_with_head_into_base_model and has_base_model_prefix:400 flax_key_tuple = flax_key_tuple[1:]401 elif load_base_model_into_model_with_head and require_base_model_prefix:402 flax_key_tuple = (pt_model.base_model_prefix,) + flax_key_tuple403 404 # rename flax weights to PyTorch format405 if flax_key_tuple[-1] == "kernel" and flax_tensor.ndim == 4 and ".".join(flax_key_tuple) not in pt_model_dict:406 # conv layer407 flax_key_tuple = flax_key_tuple[:-1] + ("weight",)408 flax_tensor = jnp.transpose(flax_tensor, (3, 2, 0, 1))409 elif flax_key_tuple[-1] == "kernel" and ".".join(flax_key_tuple) not in pt_model_dict:410 # linear layer411 flax_key_tuple = flax_key_tuple[:-1] + ("weight",)412 flax_tensor = flax_tensor.T413 elif flax_key_tuple[-1] in ["scale", "embedding"]:414 flax_key_tuple = flax_key_tuple[:-1] + ("weight",)415 416 # adding batch stats from flax batch norm to pt417 elif "mean" in flax_key_tuple[-1]:418 flax_key_tuple = flax_key_tuple[:-1] + ("running_mean",)419 elif "var" in flax_key_tuple[-1]:420 flax_key_tuple = flax_key_tuple[:-1] + ("running_var",)421 422 if "batch_stats" in flax_state:423 flax_key = ".".join(flax_key_tuple[1:]) # Remove the params/batch_stats header424 else:425 flax_key = ".".join(flax_key_tuple)426 427 # We also need to look at `pt_model_dict` and see if there are keys requiring further transformation.428 special_pt_names = {}429 # New `weight_norm` from https://github.com/huggingface/transformers/pull/24030430 for key in pt_model_dict:431 key_components = key.split(".")432 name = None433 if key_components[-3::2] == ["parametrizations", "original0"]:434 name = key_components[-2] + "_g"435 elif key_components[-3::2] == ["parametrizations", "original1"]:436 name = key_components[-2] + "_v"437 if name is not None:438 key_components = key_components[:-3] + [name]439 key_to_check = ".".join(key_components)440 special_pt_names[key_to_check] = key441 442 if flax_key in special_pt_names:443 flax_key = special_pt_names[flax_key]444 445 if flax_key in pt_model_dict:446 if flax_tensor.shape != pt_model_dict[flax_key].shape:447 raise ValueError(448 f"Flax checkpoint seems to be incorrect. Weight {flax_key_tuple} was expected "449 f"to be of shape {pt_model_dict[flax_key].shape}, but is {flax_tensor.shape}."450 )451 else:452 # add weight to pytorch dict453 flax_tensor = np.asarray(flax_tensor) if not isinstance(flax_tensor, np.ndarray) else flax_tensor454 pt_model_dict[flax_key] = torch.from_numpy(flax_tensor)455 # remove from missing keys456 missing_keys.remove(flax_key)457 else:458 # weight is not expected by PyTorch model459 unexpected_keys.append(flax_key)460 461 pt_model.load_state_dict(pt_model_dict)462 463 # re-transform missing_keys to list464 missing_keys = list(missing_keys)465 466 if len(unexpected_keys) > 0:467 logger.warning(468 "Some weights of the Flax model were not used when initializing the PyTorch model"469 f" {pt_model.__class__.__name__}: {unexpected_keys}\n- This IS expected if you are initializing"470 f" {pt_model.__class__.__name__} from a Flax model trained on another task or with another architecture"471 " (e.g. initializing a BertForSequenceClassification model from a FlaxBertForPreTraining model).\n- This"472 f" IS NOT expected if you are initializing {pt_model.__class__.__name__} from a Flax model that you expect"473 " to be exactly identical (e.g. initializing a BertForSequenceClassification model from a"474 " FlaxBertForSequenceClassification model)."475 )476 else:477 logger.warning(f"All Flax model weights were used when initializing {pt_model.__class__.__name__}.\n")478 if len(missing_keys) > 0:479 logger.warning(480 f"Some weights of {pt_model.__class__.__name__} were not initialized from the Flax model and are newly"481 f" initialized: {missing_keys}\nYou should probably TRAIN this model on a down-stream task to be able to"482 " use it for predictions and inference."483 )484 else:485 logger.warning(486 f"All the weights of {pt_model.__class__.__name__} were initialized from the Flax model.\n"487 "If your task is similar to the task the model of the checkpoint was trained on, "488 f"you can already use {pt_model.__class__.__name__} for predictions without further training."489 )490 491 return pt_model492