Convert applicable model parameters to fp32
(model: nn.Module)
| 7 | |
| 8 | |
| 9 | def convert_weights(model: nn.Module): |
| 10 | """Convert applicable model parameters to fp32""" |
| 11 | |
| 12 | def _convert_weights_to_fp32(m): |
| 13 | if isinstance(m, (nn.Conv1d, nn.Conv2d, nn.Linear)): |
| 14 | m.weight.data = m.weight.data.float() |
| 15 | if m.bias is not None: |
| 16 | m.bias.data = m.bias.data.float() |
| 17 | |
| 18 | if isinstance(m, nn.MultiheadAttention): |
| 19 | attr_list = [f"{s}_proj_weight" for s in ["in", "q", "k", "v"]] |
| 20 | attr_list += ["in_proj_bias", "bias_k", "bias_v"] |
| 21 | for attr in attr_list: |
| 22 | tensor = getattr(m, attr) |
| 23 | if tensor is not None: |
| 24 | tensor.data = tensor.data.float() |
| 25 | |
| 26 | for name in ["text_projection", "proj"]: |
| 27 | if hasattr(m, name): |
| 28 | attr = getattr(m, name) |
| 29 | if attr is not None: |
| 30 | attr.data = attr.data.float() |
| 31 | |
| 32 | model.apply(_convert_weights_to_fp32) |
| 33 | |
| 34 | |
| 35 | @SUBMODULES.register_module() |
no outgoing calls
no test coverage detected