Convert applicable model parameters to fp16
(model: nn.Module)
| 373 | |
| 374 | |
| 375 | def convert_weights(model: nn.Module): |
| 376 | """Convert applicable model parameters to fp16""" |
| 377 | |
| 378 | def _convert_weights_to_fp16(l): |
| 379 | if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Linear)): |
| 380 | l.weight.data = l.weight.data.half() |
| 381 | if l.bias is not None: |
| 382 | l.bias.data = l.bias.data.half() |
| 383 | |
| 384 | if isinstance(l, nn.MultiheadAttention): |
| 385 | for attr in [*[f"{s}_proj_weight" for s in ["in", "q", "k", "v"]], "in_proj_bias", "bias_k", "bias_v"]: |
| 386 | tensor = getattr(l, attr) |
| 387 | if tensor is not None: |
| 388 | tensor.data = tensor.data.half() |
| 389 | |
| 390 | for name in ["text_projection", "proj"]: |
| 391 | if hasattr(l, name): |
| 392 | attr = getattr(l, name) |
| 393 | if attr is not None: |
| 394 | attr.data = attr.data.half() |
| 395 | |
| 396 | model.apply(_convert_weights_to_fp16) |
| 397 | |
| 398 | |
| 399 | def build_model(state_dict: dict): |