591 lines
19 KiB
Python
591 lines
19 KiB
Python
import math
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
|
|
class TemporalAttention(nn.Module):
|
|
def __init__(
|
|
self,
|
|
n_embd: int,
|
|
n_head: int,
|
|
dropout: float = 0.0,
|
|
):
|
|
super().__init__()
|
|
assert n_embd % n_head == 0, "n_embd must be divisible by n_head"
|
|
self.n_head = n_head
|
|
self.d_head = n_embd // n_head
|
|
self.scale = 1.0 / math.sqrt(self.d_head)
|
|
|
|
# QKV projection (fused for efficiency)
|
|
self.qkv = nn.Linear(n_embd, 3 * n_embd, bias=False)
|
|
# Output projection
|
|
self.out_proj = nn.Linear(n_embd, n_embd, bias=False)
|
|
|
|
self.resid_drop = nn.Dropout(dropout)
|
|
self.reset_parameters()
|
|
|
|
def reset_parameters(self) -> None:
|
|
"""Match the previous version's GPT-style weight initialization."""
|
|
nn.init.normal_(self.qkv.weight, mean=0.0, std=0.02)
|
|
nn.init.normal_(self.out_proj.weight, mean=0.0, std=0.02)
|
|
|
|
def forward(
|
|
self,
|
|
x: torch.Tensor,
|
|
attn_mask: torch.Tensor | None = None,
|
|
) -> torch.Tensor:
|
|
B, L, _ = x.shape
|
|
H, D = self.n_head, self.d_head
|
|
|
|
# --- QKV ----------------------------------------------------------
|
|
qkv = self.qkv(x).reshape(B, L, 3, H, D).permute(2, 0, 3, 1, 4)
|
|
q, k, v = qkv.unbind(0) # each (B, H, L, D)
|
|
|
|
out = F.scaled_dot_product_attention(
|
|
q,
|
|
k,
|
|
v,
|
|
attn_mask=attn_mask,
|
|
dropout_p=0.0,
|
|
is_causal=False,
|
|
scale=self.scale,
|
|
)
|
|
|
|
# --- Aggregate & project out --------------------------------------
|
|
out = out.transpose(1, 2).reshape(B, L, H * D)
|
|
return self.resid_drop(self.out_proj(out))
|
|
|
|
|
|
class SwiGLU(nn.Module):
|
|
def __init__(
|
|
self,
|
|
n_embd: int,
|
|
hidden_dim: int | None = None,
|
|
dropout: float = 0.0,
|
|
bias: bool = True,
|
|
):
|
|
super().__init__()
|
|
hidden_dim = hidden_dim if hidden_dim is not None else int(
|
|
n_embd * 2.5)
|
|
|
|
self.w1 = nn.Linear(n_embd, hidden_dim, bias=bias) # gate path
|
|
self.w2 = nn.Linear(n_embd, hidden_dim, bias=bias) # value path
|
|
# output projection
|
|
self.w3 = nn.Linear(hidden_dim, n_embd, bias=bias)
|
|
self.drop = nn.Dropout(dropout)
|
|
self.reset_parameters()
|
|
|
|
def reset_parameters(self) -> None:
|
|
"""GPT-style parameter initialization for MLP paths."""
|
|
nn.init.normal_(self.w1.weight, mean=0.0, std=0.02)
|
|
nn.init.normal_(self.w2.weight, mean=0.0, std=0.02)
|
|
nn.init.normal_(self.w3.weight, mean=0.0, std=0.02)
|
|
if self.w1.bias is not None:
|
|
nn.init.zeros_(self.w1.bias)
|
|
nn.init.zeros_(self.w2.bias)
|
|
nn.init.zeros_(self.w3.bias)
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
"""``(B, L, n_embd) -> (B, L, n_embd)``."""
|
|
return self.drop(self.w3(F.silu(self.w1(x)) * self.w2(x)))
|
|
|
|
|
|
class GPTBlock(nn.Module):
|
|
def __init__(
|
|
self,
|
|
n_embd: int,
|
|
n_head: int,
|
|
|
|
attn_dropout: float = 0.0,
|
|
mlp_dropout: float = 0.0,
|
|
):
|
|
super().__init__()
|
|
self.attn = TemporalAttention(
|
|
n_embd=n_embd,
|
|
n_head=n_head,
|
|
dropout=attn_dropout,
|
|
)
|
|
self.mlp = SwiGLU(n_embd=n_embd, dropout=mlp_dropout)
|
|
self.ln1 = nn.LayerNorm(n_embd)
|
|
self.ln2 = nn.LayerNorm(n_embd)
|
|
|
|
def forward(
|
|
self,
|
|
x: torch.Tensor,
|
|
attn_mask: torch.Tensor | None = None,
|
|
) -> torch.Tensor:
|
|
x = x + self.attn(self.ln1(x), attn_mask)
|
|
x = x + self.mlp(self.ln2(x))
|
|
return x
|
|
|
|
|
|
class AgeSinusoidalEncoding(nn.Module):
|
|
|
|
def __init__(self, embedding_dim: int):
|
|
|
|
super().__init__()
|
|
if embedding_dim % 2 != 0:
|
|
raise ValueError(
|
|
f"Embedding dimension must be an even number, but got {embedding_dim}")
|
|
|
|
self.embedding_dim = embedding_dim
|
|
|
|
i = torch.arange(0, self.embedding_dim, 2, dtype=torch.float32)
|
|
divisor = torch.pow(10000, i / self.embedding_dim)
|
|
self.register_buffer('divisor', divisor)
|
|
self.linear = nn.Linear(embedding_dim, embedding_dim, bias=False)
|
|
|
|
def forward(self, t: torch.Tensor) -> torch.Tensor:
|
|
|
|
t_years = t
|
|
# Broadcast (B, L, 1) against (1, 1, D/2) to get (B, L, D/2)
|
|
args = t_years.unsqueeze(-1) / self.divisor.view(1, 1, -1)
|
|
# Interleave cos and sin along the last dimension
|
|
output = torch.zeros(t.shape[0], t.shape[1],
|
|
self.embedding_dim, device=t.device)
|
|
output[:, :, 0::2] = torch.cos(args)
|
|
output[:, :, 1::2] = torch.sin(args)
|
|
output = self.linear(output)
|
|
return output
|
|
|
|
|
|
class GRN2d(nn.Module):
|
|
"""Global Response Normalization from ConvNeXt V2 for NCHW tensors.
|
|
|
|
Reference:
|
|
Woo et al., "ConvNeXt V2: Co-designing and Scaling ConvNets with
|
|
Masked Autoencoders", CVPR 2023. https://arxiv.org/abs/2301.00808
|
|
"""
|
|
|
|
def __init__(self, dim: int, eps: float = 1e-6):
|
|
super().__init__()
|
|
self.gamma = nn.Parameter(torch.zeros(1, dim, 1, 1))
|
|
self.beta = nn.Parameter(torch.zeros(1, dim, 1, 1))
|
|
self.eps = eps
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
gx = torch.norm(x, p=2, dim=(2, 3), keepdim=True)
|
|
nx = gx / (gx.mean(dim=1, keepdim=True) + self.eps)
|
|
return x + self.gamma * (x * nx) + self.beta
|
|
|
|
|
|
class LayerNorm2d(nn.Module):
|
|
"""Channel-wise LayerNorm for NCHW tensors."""
|
|
|
|
def __init__(self, dim: int, eps: float = 1e-6):
|
|
super().__init__()
|
|
self.weight = nn.Parameter(torch.ones(dim))
|
|
self.bias = nn.Parameter(torch.zeros(dim))
|
|
self.eps = eps
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
x = x.permute(0, 2, 3, 1)
|
|
x = F.layer_norm(x, (self.weight.numel(),), self.weight, self.bias, self.eps)
|
|
return x.permute(0, 3, 1, 2)
|
|
|
|
|
|
class ConvNeXtV2Block2d(nn.Module):
|
|
"""Lightweight ConvNeXt V2-style 2D block for TimesNet period images.
|
|
|
|
This is intentionally a block, not the full ConvNeXt V2 image backbone:
|
|
TimesNet's 2D tensors are reshaped time-series period maps rather than
|
|
natural images, so aggressive visual downsampling would destroy axis
|
|
semantics.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
dim: int,
|
|
kernel_size: int = 7,
|
|
mlp_ratio: float = 4.0,
|
|
dropout: float = 0.0,
|
|
):
|
|
super().__init__()
|
|
padding = kernel_size // 2
|
|
hidden_dim = int(dim * mlp_ratio)
|
|
self.dwconv = nn.Conv2d(
|
|
dim,
|
|
dim,
|
|
kernel_size=kernel_size,
|
|
padding=padding,
|
|
groups=dim,
|
|
)
|
|
self.norm = LayerNorm2d(dim)
|
|
self.pwconv1 = nn.Conv2d(dim, hidden_dim, kernel_size=1)
|
|
self.act = nn.GELU()
|
|
self.grn = GRN2d(hidden_dim)
|
|
self.pwconv2 = nn.Conv2d(hidden_dim, dim, kernel_size=1)
|
|
self.drop = nn.Dropout(dropout)
|
|
self.reset_parameters()
|
|
|
|
def reset_parameters(self) -> None:
|
|
nn.init.normal_(self.dwconv.weight, mean=0.0, std=0.02)
|
|
nn.init.zeros_(self.dwconv.bias)
|
|
nn.init.normal_(self.pwconv1.weight, mean=0.0, std=0.02)
|
|
nn.init.zeros_(self.pwconv1.bias)
|
|
nn.init.normal_(self.pwconv2.weight, mean=0.0, std=0.02)
|
|
nn.init.zeros_(self.pwconv2.bias)
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
residual = x
|
|
x = self.dwconv(x)
|
|
x = self.norm(x)
|
|
x = self.pwconv1(x)
|
|
x = self.act(x)
|
|
x = self.grn(x)
|
|
x = self.pwconv2(x)
|
|
x = self.drop(x)
|
|
return residual + x
|
|
|
|
|
|
class TimesNetBlock(nn.Module):
|
|
"""TimesNet block with ConvNeXt V2-style 2D extraction.
|
|
|
|
The block follows TimesNet's idea: discover dominant periods with FFT,
|
|
reshape a 1D sequence into period-wise 2D maps, run a 2D convolutional
|
|
extractor, then fuse the top-k period branches.
|
|
|
|
Reference:
|
|
Wu et al., "TimesNet: Temporal 2D-Variation Modeling for General Time
|
|
Series Analysis", ICLR 2023. https://arxiv.org/abs/2210.02186
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
d_model: int,
|
|
top_k: int = 3,
|
|
n_convnext_blocks: int = 2,
|
|
conv_kernel_size: int = 7,
|
|
mlp_ratio: float = 4.0,
|
|
dropout: float = 0.0,
|
|
):
|
|
super().__init__()
|
|
if top_k <= 0:
|
|
raise ValueError(f"top_k must be > 0, got {top_k}")
|
|
self.top_k = top_k
|
|
self.norm = nn.LayerNorm(d_model)
|
|
self.extractor = nn.Sequential(*[
|
|
ConvNeXtV2Block2d(
|
|
dim=d_model,
|
|
kernel_size=conv_kernel_size,
|
|
mlp_ratio=mlp_ratio,
|
|
dropout=dropout,
|
|
)
|
|
for _ in range(n_convnext_blocks)
|
|
])
|
|
|
|
def _select_periods(self, x: torch.Tensor) -> tuple[list[int], torch.Tensor]:
|
|
B, T, C = x.shape
|
|
spectrum = torch.fft.rfft(x.float(), dim=1)
|
|
amplitude = spectrum.abs().mean(dim=(0, 2))
|
|
if amplitude.numel() <= 1:
|
|
return [max(T, 1)], x.new_ones(1)
|
|
|
|
amplitude = amplitude.clone()
|
|
amplitude[0] = 0.0
|
|
k = min(self.top_k, amplitude.numel() - 1)
|
|
weights, indices = torch.topk(amplitude, k=k)
|
|
periods = [max(1, T // int(idx.item())) for idx in indices]
|
|
return periods, weights.to(dtype=x.dtype, device=x.device)
|
|
|
|
def _period_branch(self, x: torch.Tensor, period: int) -> torch.Tensor:
|
|
B, T, C = x.shape
|
|
if T % period != 0:
|
|
padded_len = ((T // period) + 1) * period
|
|
pad = x.new_zeros(B, padded_len - T, C)
|
|
x_pad = torch.cat([x, pad], dim=1)
|
|
else:
|
|
padded_len = T
|
|
x_pad = x
|
|
|
|
n_periods = padded_len // period
|
|
x_2d = x_pad.reshape(B, n_periods, period, C).permute(0, 3, 1, 2)
|
|
y = self.extractor(x_2d)
|
|
y = y.permute(0, 2, 3, 1).reshape(B, padded_len, C)
|
|
return y[:, :T, :]
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
residual = x
|
|
x_norm = self.norm(x)
|
|
periods, weights = self._select_periods(x_norm)
|
|
branches = torch.stack(
|
|
[self._period_branch(x_norm, period) for period in periods],
|
|
dim=-1,
|
|
)
|
|
weights = torch.softmax(weights, dim=0).view(1, 1, 1, -1)
|
|
return residual + (branches * weights).sum(dim=-1)
|
|
|
|
|
|
class TimesNetEncoder(nn.Module):
|
|
"""Encode a multivariate time series into one fixed-size embedding."""
|
|
|
|
def __init__(
|
|
self,
|
|
input_dim: int,
|
|
d_model: int,
|
|
n_layers: int = 2,
|
|
top_k: int = 3,
|
|
n_convnext_blocks: int = 2,
|
|
conv_kernel_size: int = 7,
|
|
mlp_ratio: float = 4.0,
|
|
dropout: float = 0.0,
|
|
append_observed_mask: bool = True,
|
|
):
|
|
super().__init__()
|
|
self.input_dim = input_dim
|
|
self.append_observed_mask = append_observed_mask
|
|
in_dim = input_dim * 2 if append_observed_mask else input_dim
|
|
self.input_proj = nn.Linear(in_dim, d_model)
|
|
self.blocks = nn.ModuleList([
|
|
TimesNetBlock(
|
|
d_model=d_model,
|
|
top_k=top_k,
|
|
n_convnext_blocks=n_convnext_blocks,
|
|
conv_kernel_size=conv_kernel_size,
|
|
mlp_ratio=mlp_ratio,
|
|
dropout=dropout,
|
|
)
|
|
for _ in range(n_layers)
|
|
])
|
|
self.final_ln = nn.LayerNorm(d_model)
|
|
self.reset_parameters()
|
|
|
|
def reset_parameters(self) -> None:
|
|
nn.init.normal_(self.input_proj.weight, mean=0.0, std=0.02)
|
|
nn.init.zeros_(self.input_proj.bias)
|
|
|
|
def forward(
|
|
self,
|
|
x: torch.Tensor,
|
|
observed_mask: torch.Tensor | None = None,
|
|
) -> torch.Tensor:
|
|
if x.dim() != 3:
|
|
raise ValueError(f"x must have shape (B, T, C), got {tuple(x.shape)}")
|
|
if x.size(-1) != self.input_dim:
|
|
raise ValueError(
|
|
f"last dim must be input_dim={self.input_dim}, got {x.size(-1)}"
|
|
)
|
|
|
|
finite_mask = torch.isfinite(x)
|
|
x = torch.nan_to_num(x, nan=0.0, posinf=0.0, neginf=0.0)
|
|
if observed_mask is None:
|
|
observed_mask = finite_mask
|
|
elif observed_mask.shape == x.shape[:2]:
|
|
observed_mask = observed_mask.unsqueeze(-1).expand_as(x)
|
|
elif observed_mask.shape != x.shape:
|
|
raise ValueError(
|
|
"observed_mask must have shape (B, T) or (B, T, C), got "
|
|
f"{tuple(observed_mask.shape)}"
|
|
)
|
|
|
|
observed_mask = observed_mask.to(device=x.device, dtype=x.dtype)
|
|
x = x * observed_mask
|
|
if self.append_observed_mask:
|
|
x = torch.cat([x, observed_mask], dim=-1)
|
|
|
|
h = self.input_proj(x)
|
|
for block in self.blocks:
|
|
h = block(h)
|
|
h = self.final_ln(h)
|
|
|
|
valid_time = observed_mask.amax(dim=-1)
|
|
pooled = (h * valid_time.unsqueeze(-1)).sum(dim=1)
|
|
denom = valid_time.sum(dim=1, keepdim=True).clamp_min(1.0)
|
|
return pooled / denom
|
|
|
|
|
|
class TimesNetExposureEncoder(nn.Module):
|
|
"""Encode pre-onset environmental exposure into an event-level embedding.
|
|
|
|
Expected inputs:
|
|
daily: (B, 1826, 4) for tmean, tmax, tmin, rhmean
|
|
monthly: (B, 241, 2) for tmean, rhmean
|
|
|
|
Output:
|
|
(B, n_embd), suitable for adding to a disease event embedding in the
|
|
route-2 single-stream event-enhancement setup.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
n_embd: int,
|
|
daily_input_dim: int = 4,
|
|
monthly_input_dim: int = 2,
|
|
d_model: int | None = None,
|
|
n_layers: int = 2,
|
|
top_k: int = 3,
|
|
n_convnext_blocks: int = 2,
|
|
conv_kernel_size: int = 7,
|
|
mlp_ratio: float = 4.0,
|
|
dropout: float = 0.0,
|
|
use_gate: bool = True,
|
|
):
|
|
super().__init__()
|
|
d_model = n_embd if d_model is None else d_model
|
|
self.daily_encoder = TimesNetEncoder(
|
|
input_dim=daily_input_dim,
|
|
d_model=d_model,
|
|
n_layers=n_layers,
|
|
top_k=top_k,
|
|
n_convnext_blocks=n_convnext_blocks,
|
|
conv_kernel_size=conv_kernel_size,
|
|
mlp_ratio=mlp_ratio,
|
|
dropout=dropout,
|
|
append_observed_mask=True,
|
|
)
|
|
self.monthly_encoder = TimesNetEncoder(
|
|
input_dim=monthly_input_dim,
|
|
d_model=d_model,
|
|
n_layers=n_layers,
|
|
top_k=top_k,
|
|
n_convnext_blocks=n_convnext_blocks,
|
|
conv_kernel_size=conv_kernel_size,
|
|
mlp_ratio=mlp_ratio,
|
|
dropout=dropout,
|
|
append_observed_mask=True,
|
|
)
|
|
self.out_proj = nn.Sequential(
|
|
nn.LayerNorm(2 * d_model),
|
|
nn.Linear(2 * d_model, n_embd),
|
|
nn.GELU(),
|
|
nn.Linear(n_embd, n_embd),
|
|
)
|
|
self.gate = nn.Parameter(torch.tensor(-2.0)) if use_gate else None
|
|
self.reset_parameters()
|
|
|
|
def reset_parameters(self) -> None:
|
|
for module in self.out_proj:
|
|
if isinstance(module, nn.Linear):
|
|
nn.init.normal_(module.weight, mean=0.0, std=0.02)
|
|
nn.init.zeros_(module.bias)
|
|
|
|
def forward(
|
|
self,
|
|
daily: torch.Tensor,
|
|
monthly: torch.Tensor,
|
|
daily_mask: torch.Tensor | None = None,
|
|
monthly_mask: torch.Tensor | None = None,
|
|
) -> torch.Tensor:
|
|
h_daily = self.daily_encoder(daily, observed_mask=daily_mask)
|
|
h_monthly = self.monthly_encoder(monthly, observed_mask=monthly_mask)
|
|
h = self.out_proj(torch.cat([h_daily, h_monthly], dim=-1))
|
|
if self.gate is not None:
|
|
h = torch.sigmoid(self.gate) * h
|
|
return h
|
|
|
|
|
|
class TimesNetSequenceDecoder(nn.Module):
|
|
"""Decode a fixed-size latent vector into a multivariate time series."""
|
|
|
|
def __init__(
|
|
self,
|
|
output_dim: int,
|
|
latent_dim: int,
|
|
d_model: int,
|
|
n_layers: int = 2,
|
|
top_k: int = 3,
|
|
n_convnext_blocks: int = 2,
|
|
conv_kernel_size: int = 7,
|
|
mlp_ratio: float = 4.0,
|
|
dropout: float = 0.0,
|
|
):
|
|
super().__init__()
|
|
self.latent_proj = nn.Linear(latent_dim, d_model)
|
|
self.position_proj = nn.Linear(3, d_model)
|
|
self.blocks = nn.ModuleList([
|
|
TimesNetBlock(
|
|
d_model=d_model,
|
|
top_k=top_k,
|
|
n_convnext_blocks=n_convnext_blocks,
|
|
conv_kernel_size=conv_kernel_size,
|
|
mlp_ratio=mlp_ratio,
|
|
dropout=dropout,
|
|
)
|
|
for _ in range(n_layers)
|
|
])
|
|
self.final_ln = nn.LayerNorm(d_model)
|
|
self.output_proj = nn.Linear(d_model, output_dim)
|
|
|
|
def forward(self, latent: torch.Tensor, length: int) -> torch.Tensor:
|
|
if latent.dim() != 2:
|
|
raise ValueError(
|
|
f"latent must have shape (B, D), got {tuple(latent.shape)}"
|
|
)
|
|
position = torch.linspace(
|
|
0.0, 1.0, length, device=latent.device, dtype=latent.dtype
|
|
)
|
|
position = torch.stack(
|
|
[position, torch.sin(2 * torch.pi * position),
|
|
torch.cos(2 * torch.pi * position)],
|
|
dim=-1,
|
|
)
|
|
h = self.latent_proj(latent).unsqueeze(1)
|
|
h = h + self.position_proj(position).unsqueeze(0)
|
|
for block in self.blocks:
|
|
h = block(h)
|
|
return self.output_proj(self.final_ln(h))
|
|
|
|
|
|
class TimesNetExposureAutoencoder(nn.Module):
|
|
"""Dual-resolution exposure autoencoder with a reusable event encoder."""
|
|
|
|
def __init__(
|
|
self,
|
|
n_embd: int = 120,
|
|
daily_input_dim: int = 4,
|
|
monthly_input_dim: int = 2,
|
|
d_model: int | None = None,
|
|
n_layers: int = 2,
|
|
top_k: int = 3,
|
|
n_convnext_blocks: int = 2,
|
|
conv_kernel_size: int = 7,
|
|
mlp_ratio: float = 4.0,
|
|
dropout: float = 0.0,
|
|
):
|
|
super().__init__()
|
|
d_model = n_embd if d_model is None else d_model
|
|
encoder_kwargs = dict(
|
|
n_embd=n_embd,
|
|
daily_input_dim=daily_input_dim,
|
|
monthly_input_dim=monthly_input_dim,
|
|
d_model=d_model,
|
|
n_layers=n_layers,
|
|
top_k=top_k,
|
|
n_convnext_blocks=n_convnext_blocks,
|
|
conv_kernel_size=conv_kernel_size,
|
|
mlp_ratio=mlp_ratio,
|
|
dropout=dropout,
|
|
use_gate=True,
|
|
)
|
|
decoder_kwargs = dict(
|
|
latent_dim=n_embd,
|
|
d_model=d_model,
|
|
n_layers=n_layers,
|
|
top_k=top_k,
|
|
n_convnext_blocks=n_convnext_blocks,
|
|
conv_kernel_size=conv_kernel_size,
|
|
mlp_ratio=mlp_ratio,
|
|
dropout=dropout,
|
|
)
|
|
self.encoder = TimesNetExposureEncoder(**encoder_kwargs)
|
|
self.daily_decoder = TimesNetSequenceDecoder(
|
|
output_dim=daily_input_dim, **decoder_kwargs
|
|
)
|
|
self.monthly_decoder = TimesNetSequenceDecoder(
|
|
output_dim=monthly_input_dim, **decoder_kwargs
|
|
)
|
|
|
|
def forward(
|
|
self,
|
|
daily: torch.Tensor,
|
|
monthly: torch.Tensor,
|
|
daily_mask: torch.Tensor | None = None,
|
|
monthly_mask: torch.Tensor | None = None,
|
|
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
|
latent = self.encoder(daily, monthly, daily_mask, monthly_mask)
|
|
daily_reconstruction = self.daily_decoder(latent, daily.size(1))
|
|
monthly_reconstruction = self.monthly_decoder(latent, monthly.size(1))
|
|
return daily_reconstruction, monthly_reconstruction, latent
|