Reverts ActLinear modules back to their original nn.Linear layers.
(model)
| 98 | |
| 99 | |
| 100 | def revert_Act_to_Linear(model): |
| 101 | """ |
| 102 | Reverts ActLinear modules back to their original nn.Linear layers. |
| 103 | """ |
| 104 | for name, module in model.named_modules(): |
| 105 | if isinstance(module, ActLinear): |
| 106 | # Extract the base nn.Linear module from ActLinear |
| 107 | linear_module = module.base |
| 108 | # Navigate to the parent module of the ActLinear module |
| 109 | parent_name = name.rsplit(".", 1)[0] if "." in name else "" |
| 110 | print(f"Reverting {name}, parent: {parent_name}") |
| 111 | parent_module = ( |
| 112 | model |
| 113 | if parent_name == "" |
| 114 | else reduce(getattr, parent_name.split("."), model) |
| 115 | ) |
| 116 | # Replace the ActLinear module with the extracted nn.Linear module |
| 117 | setattr(parent_module, name.split(".")[-1], linear_module) |
| 118 | |
| 119 | return model |
| 120 | |
| 121 | |
| 122 | def clear_act_buffer(act_model): |