File size: 7,899 Bytes
38e555d | 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 177 178 179 180 181 182 183 184 185 186 187 188 189 | """Architecture ablations for the fixed-32 PASTIS fine-tuning protocol.
All existing encoder parameter names are retained for checkpoint loading.
The validity experiment changes only handling of collate/internal padding;
the seasonal experiment adds a residual temporal-pooling branch.
"""
import torch
from torch import nn
import torch.nn.functional as F
from torch.utils.checkpoint import checkpoint
import config as C
from model import (
EVAL_TEMPORAL_CHUNK, SPATIAL_CHUNK, TEMPORAL_CHUNK,
STA, UTAE_ASPP, UTAEClassificationDual,
)
from temporal_srda import TemporalSRDABlock, TemporalSRDAEncoder
def valid_from_index_positions(pos):
"""The unmodified dataset uses 0,1,...,T-1 and pads positions with zero.
Temporal subsampling may drop acquisition zero; its first retained position is
then positive. Only an initial zero is a valid acquisition-zero token.
This cannot be reused for arbitrary calendar positions without a length mask.
"""
first_is_zero = pos[:, :1].eq(0)
first = torch.zeros_like(pos, dtype=torch.bool)
first[:, :1] = first_is_zero
valid = pos.gt(0) | first
if not valid.any(dim=1).all():
raise ValueError('Each time series must contain an observation')
return valid
class MaskedTemporalSRDABlock(TemporalSRDABlock):
def forward(self, x, valid):
n, t, d = x.shape
r = self.R
x = x * valid.unsqueeze(-1)
# Remove partial-group padding *after* LayerNorm. Learned LN biases
# would otherwise turn padded zero vectors into nonzero KV features.
z = self.pw(self.norm_kv(x)) * valid.unsqueeze(-1)
z = z.reshape(n, t // r, d)
group_valid = valid.reshape(n, t // r, r).any(dim=-1)
zn = self.norm_first(z)
z_sa, _ = self.first_sa(
zn, zn, zn, key_padding_mask=~group_valid, need_weights=False
)
kv = (z + z_sa) * group_valid.unsqueeze(-1)
q = self.q_proj(self.norm_q(x))
k, v = self.kv_proj(kv).chunk(2, dim=-1)
q = q.reshape(n, t, self.h, self.dh).transpose(1, 2)
k = k.reshape(n, -1, self.h, self.dh).transpose(1, 2)
v = v.reshape(n, -1, self.h, self.dh).transpose(1, 2)
scores = (q @ k.transpose(-2, -1)) * self.scale
scores = scores.masked_fill(~group_valid[:, None, None, :], torch.finfo(scores.dtype).min)
weights = self.attn_drop(scores.softmax(dim=-1))
attended = (weights @ v).transpose(1, 2).reshape(n, t, d)
x = (x + self.proj_drop(self.out_proj(attended))) * valid.unsqueeze(-1)
y = F.gelu(self.fc1(self.norm_mlp(x))) * valid.unsqueeze(-1)
y = y + self.dw(y.transpose(1, 2)).transpose(1, 2)
x = (x + self.fc2(y)) * valid.unsqueeze(-1)
return x
class MaskedTemporalSRDAEncoder(TemporalSRDAEncoder):
def __init__(self, d_model=256, n_heads=8, R_schedule=(4, 4, 4), **kwargs):
super().__init__(d_model=d_model, n_heads=n_heads, R_schedule=R_schedule, **kwargs)
self.blocks = nn.ModuleList(
MaskedTemporalSRDABlock(d_model, n_heads, R=r) for r in R_schedule
)
def forward(self, x, pos=None, valid=None):
if valid is None:
raise ValueError('Validity mask is required for this ablation')
if pos is not None:
x = x + pos
original_t = x.shape[1]
pad = (-original_t) % self.R_max
if pad:
x = F.pad(x, (0, 0, 0, pad))
valid = F.pad(valid, (0, pad), value=False)
x = x * valid.unsqueeze(-1)
for block in self.blocks:
x = block(x, valid)
return x[:, :original_t] if pad else x
class MaskedUTAE_ASPP(UTAE_ASPP):
def __init__(self, n_channels=10, d_model=256, n_heads=8, **kwargs):
super().__init__(n_channels=n_channels, d_model=d_model, n_heads=n_heads, **kwargs)
self.temporal_encoder = MaskedTemporalSRDAEncoder(
d_model=d_model, n_heads=n_heads, R_schedule=C.R_SCHEDULE
)
def forward(self, x, pos):
b, t, channels, h, w = x.shape
valid = valid_from_index_positions(pos)
frames = x.reshape(b * t, channels, h, w)
flat_valid = valid.flatten()
# Padding is absent from the spatial BN batch as well as temporal
# attention. This matters during training with batch size two.
active = frames[flat_valid]
encoded = self._encode_spatial(active)
d = encoded.shape[1]
x2 = encoded.new_zeros((b * t, d, h, w))
x2 = x2.index_copy(0, flat_valid.nonzero().flatten(), encoded)
pixels = x2.reshape(b, t, d, h, w).permute(0, 3, 4, 1, 2).reshape(b * h * w, t, d)
pe = self.pos_encoder(pos)
expanded_pos = pe[:, None, None].expand(b, h, w, t, d).reshape(b * h * w, t, d)
expanded_valid = valid[:, None, None].expand(b, h, w, t).reshape(b * h * w, t)
n = pixels.shape[0]
if self.training and n > TEMPORAL_CHUNK:
output = torch.cat([
checkpoint(
self.temporal_encoder, pixels[s:s + TEMPORAL_CHUNK],
expanded_pos[s:s + TEMPORAL_CHUNK],
expanded_valid[s:s + TEMPORAL_CHUNK], use_reentrant=False
) for s in range(0, n, TEMPORAL_CHUNK)
], dim=0)
elif not self.training and n > EVAL_TEMPORAL_CHUNK:
output = torch.cat([
self.temporal_encoder(
pixels[s:s + EVAL_TEMPORAL_CHUNK],
expanded_pos[s:s + EVAL_TEMPORAL_CHUNK],
expanded_valid[s:s + EVAL_TEMPORAL_CHUNK],
) for s in range(0, n, EVAL_TEMPORAL_CHUNK)
], dim=0)
else:
output = self.temporal_encoder(pixels, expanded_pos, expanded_valid)
return output.reshape(b, h, w, t, d), x2
class MaskedSTA(STA):
def forward(self, feats, valid):
scores = self.l(feats)
scores = scores.masked_fill(~valid[:, None, None, :, None], torch.finfo(scores.dtype).min)
weights = scores.softmax(dim=3)
pooled = (feats * weights).sum(dim=3).permute(0, 3, 1, 2).contiguous()
pooled = pooled + self.conv2(self.conv1(pooled))
return self.conv3(pooled)
class MaskedUTAEClassificationDual(UTAEClassificationDual):
def __init__(self, encoder, d_model=256, num_classes=20):
super().__init__(encoder, d_model=d_model, num_classes=num_classes)
self.sta = MaskedSTA(d_model)
def forward(self, x, pos):
features, _ = self.utae(x, pos)
dense = self.sta(features, valid_from_index_positions(pos))
semantic = self.sem_head(dense)
boundary = self.bnd_head(dense)
return {
'sem_logits': semantic,
'bnd_logits': boundary,
'refined_logits': self.refine(dense, semantic, boundary),
}
class SeasonalSTA(STA):
def __init__(self, d_model=256, n_queries=4):
super().__init__(d_model)
self.seasonal_scores = nn.Linear(d_model, n_queries)
self.seasonal_projection = nn.Conv2d(n_queries * d_model, d_model, 1)
nn.init.zeros_(self.seasonal_projection.weight)
nn.init.zeros_(self.seasonal_projection.bias)
def forward(self, feats):
base = super().forward(feats)
# Each query pools a different temporal distribution at every pixel.
weights = self.seasonal_scores(feats).softmax(dim=3)
pooled = torch.einsum('bhwtd,bhwtq->bhwqd', feats, weights)
b, h, w, q, d = pooled.shape
pooled = pooled.reshape(b, h, w, q * d).permute(0, 3, 1, 2).contiguous()
return base + self.seasonal_projection(pooled)
class SeasonalUTAEClassificationDual(UTAEClassificationDual):
def __init__(self, encoder, d_model=256, num_classes=20):
super().__init__(encoder, d_model=d_model, num_classes=num_classes)
self.sta = SeasonalSTA(d_model)
|