This does the final conversion step: take locally converted weights and apply a global renaming to them. It splits attention layers, and takes into account additional replacements that may arise. Assigns the weights to the new checkpoint.
(
paths, checkpoint, old_checkpoint, attention_paths_to_split=None, additional_replacements=None, config=None
)
| 26 | |
| 27 | |
| 28 | def assign_to_checkpoint( |
| 29 | paths, checkpoint, old_checkpoint, attention_paths_to_split=None, additional_replacements=None, config=None |
| 30 | ): |
| 31 | """ |
| 32 | This does the final conversion step: take locally converted weights and apply a global renaming to them. It splits |
| 33 | attention layers, and takes into account additional replacements that may arise. |
| 34 | |
| 35 | Assigns the weights to the new checkpoint. |
| 36 | """ |
| 37 | assert isinstance(paths, list), "Paths should be a list of dicts containing 'old' and 'new' keys." |
| 38 | |
| 39 | # Splits the attention layers into three variables. |
| 40 | if attention_paths_to_split is not None: |
| 41 | for path, path_map in attention_paths_to_split.items(): |
| 42 | old_tensor = old_checkpoint[path] |
| 43 | channels = old_tensor.shape[0] // 3 |
| 44 | |
| 45 | target_shape = (-1, channels) if len(old_tensor.shape) == 3 else (-1) |
| 46 | |
| 47 | num_heads = old_tensor.shape[0] // config["num_head_channels"] // 3 |
| 48 | |
| 49 | old_tensor = old_tensor.reshape((num_heads, 3 * channels // num_heads) + old_tensor.shape[1:]) |
| 50 | query, key, value = old_tensor.split(channels // num_heads, dim=1) |
| 51 | |
| 52 | checkpoint[path_map["query"]] = query.reshape(target_shape) |
| 53 | checkpoint[path_map["key"]] = key.reshape(target_shape) |
| 54 | checkpoint[path_map["value"]] = value.reshape(target_shape) |
| 55 | |
| 56 | for path in paths: |
| 57 | new_path = path["new"] |
| 58 | |
| 59 | # These have already been assigned |
| 60 | if attention_paths_to_split is not None and new_path in attention_paths_to_split: |
| 61 | continue |
| 62 | |
| 63 | if additional_replacements is not None: |
| 64 | for replacement in additional_replacements: |
| 65 | new_path = new_path.replace(replacement["old"], replacement["new"]) |
| 66 | |
| 67 | # proj_attn.weight has to be converted from conv 1D to linear |
| 68 | weight = old_checkpoint[path["old"]] |
| 69 | names = ["proj_attn.weight"] |
| 70 | names_2 = ["proj_out.weight", "proj_in.weight"] |
| 71 | if any(k in new_path for k in names): |
| 72 | checkpoint[new_path] = weight[:, :, 0] |
| 73 | elif any(k in new_path for k in names_2) and len(weight.shape) > 2 and ".attentions." not in new_path: |
| 74 | checkpoint[new_path] = weight[:, :, 0] |
| 75 | else: |
| 76 | checkpoint[new_path] = weight |
| 77 | |
| 78 | |
| 79 | def renew_attention_paths(old_list, n_shave_prefix_segments=0): |
no test coverage detected
searching dependent graphs…