TSRDA / PASTIS /main_method /code /top5_models.py

Move PASTIS work into PASTIS/ (3518/3518)

0ad6e78 verified 2 days ago

7.9 kB

"""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)

	"""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)