Spaces:
Running on Zero
Running on Zero
File size: 5,082 Bytes
751d508 | 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 | """
Patching for 1D sequences, as in PatchTST: [B, L] -> [B, num_patches, d_model].
"""
import torch
import torch.nn as nn
from typing import Optional
class Patching(nn.Module):
"""
Convert time series into patches and embed them.
Takes a time series of length L and converts it into num_patches patches,
each of length patch_size. Then projects each patch to d_model dimensions.
Args:
patch_size: Number of timesteps in each patch
d_model: Embedding dimension
num_features: Number of input features/channels (default: 1 for univariate)
stride: Stride for patching (if None, uses patch_size for non-overlapping)
padding: Whether to pad the sequence to make it divisible by patch_size
"""
def __init__(
self,
patch_size: int,
d_model: int,
num_features: int = 1,
stride: Optional[int] = None,
padding: bool = True
):
super().__init__()
self.patch_size = patch_size
self.d_model = d_model
self.num_features = num_features
self.stride = stride if stride is not None else patch_size
self.padding = padding
self.projection = nn.Linear(patch_size * num_features, d_model)
# Unused in forward; kept so existing checkpoints load unchanged.
self.value_embedding = nn.Linear(num_features, d_model)
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""
Convert time series to patches.
Args:
x: Input tensor [B, L] or [B, L, C]
Returns:
Patched tensor [B, num_patches, d_model]
"""
batch_size, seq_len = x.shape[0], x.shape[1]
if x.dim() == 2:
x = x.unsqueeze(-1) # [B, L, 1]
if self.padding:
remainder = (seq_len - self.patch_size) % self.stride
if remainder != 0:
pad_len = self.stride - remainder
# Pad at the end by repeating the last value.
x = torch.cat([x, x[:, -1:, :].repeat(1, pad_len, 1)], dim=1)
seq_len = x.shape[1]
num_patches = (seq_len - self.patch_size) // self.stride + 1
patches = x.unfold(dimension=1, size=self.patch_size, step=self.stride) # [B, num_patches, C, patch_size]
patches = patches.transpose(2, 3) # [B, num_patches, patch_size, C]
patches = patches.reshape(batch_size, num_patches, -1) # [B, num_patches, patch_size * C]
patches = self.projection(patches) # [B, num_patches, d_model]
return patches
class UnPatching(nn.Module):
"""
Convert patches back to a continuous time series.
Args:
patch_size: Size of each patch
d_model: Embedding dimension
num_features: Number of output features
stride: Stride used in patching
"""
def __init__(
self,
patch_size: int,
d_model: int,
num_features: int = 1,
stride: Optional[int] = None
):
super().__init__()
self.patch_size = patch_size
self.d_model = d_model
self.num_features = num_features
self.stride = stride if stride is not None else patch_size
self.projection = nn.Linear(d_model, patch_size * num_features)
def forward(
self,
x: torch.Tensor,
target_len: Optional[int] = None
) -> torch.Tensor:
"""
Convert patches back to continuous sequence.
Args:
x: Patched tensor [B, num_patches, d_model]
target_len: Target sequence length (if None, inferred from patches)
Returns:
Continuous sequence [B, L, C]
"""
batch_size, num_patches, _ = x.shape
x = self.projection(x) # [B, num_patches, patch_size * C]
x = x.reshape(batch_size, num_patches, self.patch_size, self.num_features)
if self.stride == self.patch_size:
x = x.reshape(batch_size, -1, self.num_features) # [B, num_patches * patch_size, C]
else:
# Overlapping patches: average the overlapping regions.
seq_len = (num_patches - 1) * self.stride + self.patch_size
output = torch.zeros(
batch_size, seq_len, self.num_features,
device=x.device, dtype=x.dtype
)
counts = torch.zeros(seq_len, device=x.device, dtype=x.dtype)
for i in range(num_patches):
start = i * self.stride
end = start + self.patch_size
output[:, start:end, :] += x[:, i, :, :]
counts[start:end] += 1
x = output / counts.unsqueeze(0).unsqueeze(-1)
if target_len is not None:
x = x[:, :target_len, :]
return x |