(encoder: nn.Module, decoder: nn.Module,
base_model_prefix: str, skip_key: str)
| 19 | |
| 20 | |
| 21 | def tie_encoder_decoder_weights(encoder: nn.Module, decoder: nn.Module, |
| 22 | base_model_prefix: str, skip_key: str): |
| 23 | uninitialized_encoder_weights: List[str] = [] |
| 24 | if decoder.__class__ != encoder.__class__: |
| 25 | logger.info( |
| 26 | f"{decoder.__class__} and {encoder.__class__} are not equal. In this case make sure that all encoder weights are correctly initialized." |
| 27 | ) |
| 28 | |
| 29 | def tie_encoder_to_decoder_recursively( |
| 30 | decoder_pointer: nn.Module, |
| 31 | encoder_pointer: nn.Module, |
| 32 | module_name: str, |
| 33 | uninitialized_encoder_weights: List[str], |
| 34 | skip_key: str, |
| 35 | depth=0, |
| 36 | ): |
| 37 | assert isinstance(decoder_pointer, nn.Module) and isinstance( |
| 38 | encoder_pointer, nn.Module |
| 39 | ), f"{decoder_pointer} and {encoder_pointer} have to be of type torch.nn.Module" |
| 40 | if hasattr(decoder_pointer, "weight") and skip_key not in module_name: |
| 41 | assert hasattr(encoder_pointer, "weight") |
| 42 | encoder_pointer.weight = decoder_pointer.weight |
| 43 | if hasattr(decoder_pointer, "bias"): |
| 44 | assert hasattr(encoder_pointer, "bias") |
| 45 | encoder_pointer.bias = decoder_pointer.bias |
| 46 | print(module_name + ' is tied') |
| 47 | return |
| 48 | |
| 49 | encoder_modules = encoder_pointer._modules |
| 50 | decoder_modules = decoder_pointer._modules |
| 51 | if len(decoder_modules) > 0: |
| 52 | assert ( |
| 53 | len(encoder_modules) > 0 |
| 54 | ), f"Encoder module {encoder_pointer} does not match decoder module {decoder_pointer}" |
| 55 | |
| 56 | all_encoder_weights = set([ |
| 57 | module_name + "/" + sub_name |
| 58 | for sub_name in encoder_modules.keys() |
| 59 | ]) |
| 60 | encoder_layer_pos = 0 |
| 61 | for name, module in decoder_modules.items(): |
| 62 | if name.isdigit(): |
| 63 | encoder_name = str(int(name) + encoder_layer_pos) |
| 64 | decoder_name = name |
| 65 | if not isinstance( |
| 66 | decoder_modules[decoder_name], |
| 67 | type(encoder_modules[encoder_name])) and len( |
| 68 | encoder_modules) != len(decoder_modules): |
| 69 | # this can happen if the name corresponds to the position in a list module list of layers |
| 70 | # in this case the decoder has added a cross-attention that the encoder does not have |
| 71 | # thus skip this step and subtract one layer pos from encoder |
| 72 | encoder_layer_pos -= 1 |
| 73 | continue |
| 74 | elif name not in encoder_modules: |
| 75 | continue |
| 76 | elif depth > 500: |
| 77 | raise ValueError( |
| 78 | "Max depth of recursive function `tie_encoder_to_decoder` reached. It seems that there is a circular dependency between two or more `nn.Modules` of your model." |
no test coverage detected