File size: 8,501 Bytes
4c4d99c | 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 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 | """Pure-PyTorch engineering reproduction of the Prithvi-EO-2.0 TL MAE."""
import math
import torch
from torch import nn
def sincos_1d(positions, dim):
if dim % 2:
raise ValueError("sine/cosine dimensions must be even")
frequencies = torch.exp(
torch.arange(dim // 2, device=positions.device, dtype=positions.dtype)
* (-math.log(10000.0) / max(dim // 2, 1))
)
angles = positions.unsqueeze(-1) * frequencies
return torch.cat((angles.sin(), angles.cos()), dim=-1)
def sincos_3d(frames, height, width, dim, device, dtype):
if dim % 16:
raise ValueError("3D position dimension must be divisible by 16")
width_dim, height_dim, time_dim = 6 * dim // 16, 6 * dim // 16, 4 * dim // 16
time, row, column = torch.meshgrid(
torch.arange(frames, device=device, dtype=dtype),
torch.arange(height, device=device, dtype=dtype),
torch.arange(width, device=device, dtype=dtype),
indexing="ij",
)
return torch.cat((
sincos_1d(column.reshape(-1), width_dim),
sincos_1d(row.reshape(-1), height_dim),
sincos_1d(time.reshape(-1), time_dim),
), dim=-1)
def patchify(values, patch_size):
batch, channels, frames, height, width = values.shape
pt, ph, pw = patch_size
if frames % pt or height % ph or width % pw:
raise ValueError("input dimensions must be divisible by patch_size")
return values.reshape(
batch, channels, frames // pt, pt, height // ph, ph, width // pw, pw
).permute(0, 2, 4, 6, 3, 5, 7, 1).reshape(batch, -1, pt * ph * pw * channels)
def unpatchify(patches, channels, output_size, patch_size):
batch = patches.shape[0]
frames, height, width = output_size
pt, ph, pw = patch_size
return patches.reshape(
batch, frames // pt, height // ph, width // pw, pt, ph, pw, channels
).permute(0, 7, 1, 4, 2, 5, 3, 6).reshape(batch, channels, frames, height, width)
class Transformer(nn.Module):
def __init__(self, dim, depth, heads, mlp_ratio):
super().__init__()
layer = nn.TransformerEncoderLayer(
dim, heads, int(dim * mlp_ratio), activation="gelu", batch_first=True, norm_first=True
)
self.blocks = nn.TransformerEncoder(layer, depth)
self.norm = nn.LayerNorm(dim)
def forward(self, values):
return self.norm(self.blocks(values))
class CoordinateEncoder(nn.Module):
def __init__(self, dim, scale=0.1):
super().__init__()
if dim % 4:
raise ValueError("coordinate embedding dimension must be divisible by four")
self.dim = dim
self.scale = nn.Parameter(torch.tensor(float(scale)))
def forward(self, coordinates):
return self.scale * torch.cat((
sincos_1d(coordinates[..., 0], self.dim // 2),
sincos_1d(coordinates[..., 1], self.dim // 2),
), dim=-1)
class PrithviEO2(nn.Module):
def __init__(self, config):
super().__init__()
self.config = dict(config)
self.input_size = tuple(int(value) for value in config["input_size"])
self.patch_size = tuple(int(value) for value in config["patch_size"])
self.channels = int(config["channels"])
self.mask_ratio = float(config["mask_ratio"])
self.metadata_dropout = float(config["metadata_dropout"])
self.norm_pix_loss = bool(config.get("norm_pix_loss", False))
enc_dim, dec_dim = int(config["encoder_dim"]), int(config["decoder_dim"])
self.patch_embed = nn.Conv3d(
self.channels, enc_dim, kernel_size=self.patch_size, stride=self.patch_size
)
self.cls_token = nn.Parameter(torch.randn(1, 1, enc_dim) * 0.02)
self.encoder = Transformer(enc_dim, int(config["encoder_depth"]), int(config["encoder_heads"]),
float(config["mlp_ratio"]))
self.encoder_to_decoder = nn.Linear(enc_dim, dec_dim)
self.mask_token = nn.Parameter(torch.randn(1, 1, dec_dim) * 0.02)
self.decoder = Transformer(dec_dim, int(config["decoder_depth"]), int(config["decoder_heads"]),
float(config["mlp_ratio"]))
patch_volume = math.prod(self.patch_size) * self.channels
self.decoder_prediction = nn.Linear(dec_dim, patch_volume)
self.time_encoder = CoordinateEncoder(enc_dim)
self.location_encoder = CoordinateEncoder(enc_dim)
self.decoder_time_encoder = CoordinateEncoder(dec_dim)
self.decoder_location_encoder = CoordinateEncoder(dec_dim)
def _grid(self, pixels):
return tuple(size // patch for size, patch in zip(pixels.shape[-3:], self.patch_size))
def _metadata(self, temporal, location, grid, encoder=True):
frames, height, width = grid
time_encoder = self.time_encoder if encoder else self.decoder_time_encoder
location_encoder = self.location_encoder if encoder else self.decoder_location_encoder
temporal_embedding = time_encoder(temporal)
temporal_embedding = temporal_embedding[:, :, None, :].expand(-1, -1, height * width, -1).reshape(
len(temporal), frames * height * width, -1
)
location_embedding = location_encoder(location)[:, None, :].expand(-1, frames * height * width, -1)
if self.training and self.metadata_dropout:
time_keep = (torch.rand(len(temporal), 1, 1, device=temporal.device) >= self.metadata_dropout).to(temporal.dtype)
location_keep = (torch.rand(len(location), 1, 1, device=location.device) >= self.metadata_dropout).to(location.dtype)
temporal_embedding = temporal_embedding * time_keep
location_embedding = location_embedding * location_keep
return temporal_embedding + location_embedding
def _encoded_tokens(self, pixels, temporal, location):
grid = self._grid(pixels)
tokens = self.patch_embed(pixels).flatten(2).transpose(1, 2)
position = sincos_3d(*grid, tokens.shape[-1], tokens.device, tokens.dtype)
tokens = tokens + position[None] + self._metadata(temporal, location, grid, encoder=True)
return tokens, grid
def encode(self, pixels, temporal, location):
tokens, _ = self._encoded_tokens(pixels, temporal, location)
cls = self.cls_token.expand(len(pixels), -1, -1)
encoded = self.encoder(torch.cat((cls, tokens), dim=1))
return encoded[:, 0], encoded[:, 1:]
def forward(self, pixels, temporal, location, mask_ratio=None):
ratio = self.mask_ratio if mask_ratio is None else float(mask_ratio)
tokens, grid = self._encoded_tokens(pixels, temporal, location)
batch, length, dim = tokens.shape
keep = max(1, int(length * (1.0 - ratio)))
ordering = torch.rand(batch, length, device=pixels.device).argsort(dim=1)
visible_indices, masked_indices = ordering[:, :keep], ordering[:, keep:]
visible = tokens.gather(1, visible_indices[:, :, None].expand(-1, -1, dim))
encoded = self.encoder(torch.cat((self.cls_token.expand(batch, -1, -1), visible), dim=1))
embedding = encoded[:, 0]
visible_decoder = self.encoder_to_decoder(encoded[:, 1:])
decoder_tokens = self.mask_token.expand(batch, length, -1).clone()
decoder_tokens.scatter_(1, visible_indices[:, :, None].expand(-1, -1, visible_decoder.shape[-1]), visible_decoder)
position = sincos_3d(*grid, decoder_tokens.shape[-1], decoder_tokens.device, decoder_tokens.dtype)
decoder_tokens = decoder_tokens + position[None] + self._metadata(temporal, location, grid, encoder=False)
predictions = self.decoder_prediction(self.decoder(decoder_tokens))
targets = patchify(pixels, self.patch_size)
if self.norm_pix_loss:
mean, variance = targets.mean(dim=-1, keepdim=True), targets.var(dim=-1, keepdim=True)
targets = (targets - mean) / (variance + 1e-6).sqrt()
mask = torch.zeros(batch, length, device=pixels.device)
mask.scatter_(1, masked_indices, 1.0)
patch_mse = (predictions - targets).pow(2).mean(dim=-1)
loss = (patch_mse * mask).sum() / mask.sum().clamp_min(1)
reconstruction = unpatchify(predictions, self.channels, pixels.shape[-3:], self.patch_size)
return {
"loss": loss,
"embedding": embedding,
"patch_embeddings": encoded[:, 1:],
"reconstruction": reconstruction,
"mask": mask,
}
|