Reverts ActLinear modules back to their original nn.Linear layers.
(model)
| 113 | |
| 114 | |
| 115 | def revert_Act_to_Linear(model): |
| 116 | """ |
| 117 | Reverts ActLinear modules back to their original nn.Linear layers. |
| 118 | """ |
| 119 | for name, module in model.named_modules(): |
| 120 | if isinstance(module, ActLinear): |
| 121 | # Extract the base nn.Linear module from ActLinear |
| 122 | linear_module = module.base |
| 123 | # Navigate to the parent module of the ActLinear module |
| 124 | parent_name = name.rsplit(".", 1)[0] if "." in name else "" |
| 125 | print(f"Reverting {name}, parent: {parent_name}") |
| 126 | parent_module = ( |
| 127 | model |
| 128 | if parent_name == "" |
| 129 | else reduce(getattr, parent_name.split("."), model) |
| 130 | ) |
| 131 | # Replace the ActLinear module with the extracted nn.Linear module |
| 132 | setattr(parent_module, name.split(".")[-1], linear_module) |
| 133 | |
| 134 | return model |
| 135 | |
| 136 | |
| 137 | def clear_act_buffer(act_model): |
no outgoing calls
no test coverage detected