| """Copyright (c) Microsoft Corporation. Licensed under the MIT license.""" |
|
|
| import math |
| from typing import TypeVar |
|
|
| import torch |
| from einops import rearrange |
| from timm.models.vision_transformer import trunc_normal_ |
| from torch import nn |
|
|
| __all__ = [ |
| "unpatchify", |
| "check_lat_lon_dtype", |
| "maybe_adjust_windows", |
| "init_weights", |
| "fp16_safe_scaled_dot_product_attention", |
| ] |
|
|
|
|
| def unpatchify(x: torch.Tensor, V: int, H: int, W: int, P: int) -> torch.Tensor: |
| """Unpatchify hidden representation. |
| |
| Args: |
| x (torch.Tensor): Patchified input of shape `(B, L, C, V * P^2)` where `P` is the |
| patch size. |
| V (int): Number of variables. |
| H (int): Number of latitudes. |
| W (int): Number of longitudes. |
| P (int): Patch size. |
| |
| Returns: |
| torch.Tensor: Unpatchified representation of shape `(B, V, C, H, W)`. |
| """ |
| assert x.dim() == 4, f"Expected 4D tensor, but got {x.dim()}D." |
| B, C = x.size(0), x.size(2) |
| H = H // P |
| W = W // P |
| assert x.size(1) == H * W |
| assert x.size(-1) == V * P**2 |
|
|
| x = x.reshape(shape=(B, H, W, C, P, P, V)) |
| x = rearrange(x, "B H W C P1 P2 V -> B V C H P1 W P2") |
| x = x.reshape(shape=(B, V, C, H * P, W * P)) |
| return x |
|
|
|
|
| def check_lat_lon_dtype(lat: torch.Tensor, lon: torch.Tensor) -> None: |
| """Assert that `lat` and `lon` are at least `float32`s.""" |
| assert lat.dtype in [torch.float32, torch.float64], f"Latitude num. unstable: {lat.dtype}." |
| assert lon.dtype in [torch.float32, torch.float64], f"Longitude num. unstable: {lon.dtype}." |
|
|
|
|
| T = TypeVar("T", tuple[int, int], tuple[int, int, int]) |
|
|
|
|
| def maybe_adjust_windows(window_size: T, shift_size: T, res: T) -> tuple[T, T]: |
| """Adjust the window size and shift size if the input resolution is smaller than the window |
| size.""" |
| err_msg = f"Expected same length, found {len(window_size)}, {len(shift_size)} and {len(res)}." |
| assert len(window_size) == len(shift_size) == len(res), err_msg |
|
|
| mut_shift_size, mut_window_size = list(shift_size), list(window_size) |
| for i in range(len(res)): |
| if res[i] <= window_size[i]: |
| mut_shift_size[i] = 0 |
| mut_window_size[i] = res[i] |
|
|
| new_window_size: T = tuple(mut_window_size) |
| new_shift_size: T = tuple(mut_shift_size) |
|
|
| assert min(new_window_size) > 0, f"Window size must be positive. Found {new_window_size}." |
| assert min(new_shift_size) >= 0, f"Shift size must be non-negative. Found {new_shift_size}." |
|
|
| return new_window_size, new_shift_size |
|
|
|
|
| def fp16_safe_scaled_dot_product_attention( |
| query: torch.Tensor, |
| key: torch.Tensor, |
| value: torch.Tensor, |
| attn_mask: torch.Tensor | None = None, |
| dropout_p: float = 0.0, |
| scale: float | None = None, |
| ) -> torch.Tensor: |
| """Scaled dot-product attention with float16 overflow protection. |
| |
| Equivalent to :func:`torch.nn.functional.scaled_dot_product_attention`, but clamps intermediate |
| attention weights when running in float16 to prevent overflow or inf values that can appear with |
| large sequence lengths. |
| """ |
| scale_factor = 1 / math.sqrt(query.size(-1)) if scale is None else scale |
| |
| attn_weight = query @ (key.transpose(-2, -1) * scale_factor) |
| if attn_weight.dtype == torch.float16: |
| max_val = torch.finfo(attn_weight.dtype).max |
| clamp_value = torch.where(torch.isinf(attn_weight).any(), max_val - 1000, max_val) |
| attn_weight = torch.clamp(attn_weight, min=-clamp_value, max=clamp_value) |
| if attn_mask is not None: |
| attn_weight = attn_weight + attn_mask |
| attn_weight = torch.softmax(attn_weight, dim=-1) |
| if dropout_p > 0.0: |
| attn_weight = torch.dropout(attn_weight, dropout_p, train=True) |
| return attn_weight @ value |
|
|
|
|
| def init_weights(m: nn.Module): |
| """Initialise weights of a module with a truncated normal distribution. |
| |
| `nn.LayerNorm` is initialised with a `weight` of 1 and a `bias` of 0. |
| |
| Args: |
| m (torch.nn.Module): Module. |
| """ |
| if isinstance(m, (nn.Linear, nn.Conv2d, nn.Conv3d, nn.ConvTranspose2d, nn.ConvTranspose3d)): |
| trunc_normal_(m.weight, std=0.02) |
| if m.bias is not None: |
| nn.init.constant_(m.bias, 0) |
| elif isinstance(m, nn.LayerNorm): |
| if m.bias is not None: |
| nn.init.constant_(m.bias, 0) |
| if m.weight is not None: |
| nn.init.constant_(m.weight, 1.0) |
|
|
