File size: 4,507 Bytes
986e37d | 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 | from __future__ import annotations
import math
import torch
from torch import nn
def fourier_loglam(loglam: torch.Tensor, n_freq: int = 32) -> torch.Tensor:
# Center roughly around optical wavelengths to improve conditioning.
x = loglam - math.log(6000.0)
freqs = torch.logspace(0.0, math.log10(128.0), n_freq, device=loglam.device, dtype=loglam.dtype)
angles = x.unsqueeze(-1) * freqs
return torch.cat([torch.sin(angles), torch.cos(angles)], dim=-1)
class NativeSpecZMAE(nn.Module):
def __init__(
self,
feature_dim: int = 6,
target_aux_dim: int = 4,
d_model: int = 256,
nhead: int = 8,
encoder_layers: int = 4,
decoder_layers: int = 3,
dim_feedforward: int | None = None,
dropout: float = 0.1,
fourier_freqs: int = 32,
z_bins: int = 64,
y_min: float = 0.0,
y_max: float = math.log1p(6.0),
):
super().__init__()
self.fourier_freqs = fourier_freqs
self.y_min = y_min
self.y_max = y_max
self.z_bins = z_bins
pos_dim = fourier_freqs * 2
ff = dim_feedforward or d_model * 4
self.input_proj = nn.Sequential(
nn.Linear(feature_dim + pos_dim, d_model),
nn.LayerNorm(d_model),
nn.GELU(),
nn.Linear(d_model, d_model),
)
self.target_proj = nn.Sequential(
nn.Linear(target_aux_dim + pos_dim, d_model),
nn.LayerNorm(d_model),
nn.GELU(),
nn.Linear(d_model, d_model),
)
self.cls = nn.Parameter(torch.randn(1, 1, d_model) * 0.02)
self.z_query = nn.Parameter(torch.randn(1, 1, d_model) * 0.02)
enc_layer = nn.TransformerEncoderLayer(
d_model=d_model,
nhead=nhead,
dim_feedforward=ff,
dropout=dropout,
batch_first=True,
norm_first=True,
activation="gelu",
)
self.encoder = nn.TransformerEncoder(enc_layer, num_layers=encoder_layers)
dec_layer = nn.TransformerDecoderLayer(
d_model=d_model,
nhead=nhead,
dim_feedforward=ff,
dropout=dropout,
batch_first=True,
norm_first=True,
activation="gelu",
)
self.decoder = nn.TransformerDecoder(dec_layer, num_layers=decoder_layers)
self.rec_head = nn.Sequential(nn.LayerNorm(d_model), nn.Linear(d_model, d_model // 2), nn.GELU(), nn.Linear(d_model // 2, 1))
self.z_head = nn.Sequential(nn.LayerNorm(d_model * 2), nn.Linear(d_model * 2, d_model), nn.GELU(), nn.Linear(d_model, 2))
self.z_bin_head = nn.Sequential(nn.LayerNorm(d_model * 2), nn.Linear(d_model * 2, d_model), nn.GELU(), nn.Linear(d_model, z_bins))
def forward(
self,
enc_features: torch.Tensor,
enc_loglam: torch.Tensor,
enc_padding: torch.Tensor,
target_loglam: torch.Tensor,
target_aux: torch.Tensor,
) -> dict[str, torch.Tensor]:
bsz = enc_features.shape[0]
enc_pos = fourier_loglam(enc_loglam, self.fourier_freqs)
tokens = self.input_proj(torch.cat([enc_features, enc_pos], dim=-1))
cls = self.cls.expand(bsz, -1, -1)
z_query = self.z_query.expand(bsz, -1, -1)
memory_in = torch.cat([cls, z_query, tokens], dim=1)
special_padding = torch.zeros((bsz, 2), dtype=torch.bool, device=enc_padding.device)
memory_padding = torch.cat([special_padding, enc_padding], dim=1)
memory = self.encoder(memory_in, src_key_padding_mask=memory_padding)
cls_out = memory[:, 0]
z_out = memory[:, 1]
z_feat = torch.cat([z_out, cls_out], dim=-1)
z_params = self.z_head(z_feat)
z_bins = self.z_bin_head(z_feat)
tgt_pos = fourier_loglam(target_loglam, self.fourier_freqs)
tgt = self.target_proj(torch.cat([target_aux, tgt_pos], dim=-1))
decoded = self.decoder(tgt=tgt, memory=memory, memory_key_padding_mask=memory_padding)
rec = self.rec_head(decoded).squeeze(-1)
return {
"rec": rec,
"y_mu": z_params[:, 0],
"y_logvar": torch.clamp(z_params[:, 1], -8.0, 4.0),
"z_bin_logits": z_bins,
}
def y_to_bin(self, y: torch.Tensor) -> torch.Tensor:
scaled = (y - self.y_min) / max(self.y_max - self.y_min, 1e-6)
return torch.clamp((scaled * self.z_bins).long(), 0, self.z_bins - 1)
|