File size: 4,535 Bytes
87b732d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 | """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) # type: ignore[assignment]
new_shift_size: T = tuple(mut_shift_size) # type: ignore[assignment]
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
# Multiply scale into the key (instead of the result) to keep magnitudes lower.
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)
|