如何在PyTorch中结合Transformer实现适配特定输入形状的Time2Vec？

Time2Vec PyTorch实现（对应目标Keras版本逻辑） #

这个实现和你参考的Keras版本核心逻辑一致，原生支持处理(batch_size, seq_len, features)形状的输入：

import torch
import torch.nn as nn

class Time2Vec(nn.Module):
    def __init__(self, in_features, out_features, activation='sin'):
        super().__init__()
        self.out_features = out_features
        # 线性分量层
        self.w0 = nn.Linear(in_features, 1)
        # 周期分量层
        self.W = nn.Linear(in_features, out_features - 1)
        self.activation = torch.sin if activation == 'sin' else torch.cos

    def forward(self, x):
        # x: (batch_size, seq_len, in_features)
        linear_out = self.w0(x)  # 输出形状: (batch_size, seq_len, 1)
        periodic_out = self.activation(self.W(x))  # 输出形状: (batch_size, seq_len, out_features-1)
        # 拼接线性和周期分量
        return torch.cat([linear_out, periodic_out], dim=-1)

结合Transformer的完整示例 #

下面是适配(batch_size, seq_len, features)输入的时间序列预测模型，将Time2Vec编码后的时间特征与原特征结合后送入Transformer：

class T2VTransformerModel(nn.Module):
    def __init__(self, input_feature_num, t2v_output_dim, d_model, nhead, transformer_layers, output_dim):
        super().__init__()
        # 假设输入特征中第0维是时间特征，初始化Time2Vec层
        self.t2v_layer = Time2Vec(in_features=1, out_features=t2v_output_dim)
        # 将原特征+Time2Vec编码特征投影到Transformer所需的d_model维度
        self.feature_proj = nn.Linear(input_feature_num + t2v_output_dim, d_model)
        # Transformer编码器（batch_first=True适配(batch, seq, feature)输入）
        encoder_layer = nn.TransformerEncoderLayer(d_model=d_model, nhead=nhead, batch_first=True)
        self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=transformer_layers)
        # 最终预测层
        self.predictor = nn.Linear(d_model, output_dim)

    def forward(self, x):
        # x: (batch_size, seq_len, input_feature_num)
        # 提取时间特征（根据你的数据结构调整索引，这里取第0维）
        time_feat = x[..., 0:1]  # 形状: (batch_size, seq_len, 1)
        # 对时间特征做Time2Vec编码
        t2v_encoded = self.t2v_layer(time_feat)  # 形状: (batch_size, seq_len, t2v_output_dim)
        # 拼接原特征与编码后的时间特征
        combined_feat = torch.cat([x, t2v_encoded], dim=-1)  # 形状: (batch_size, seq_len, input_feature_num + t2v_output_dim)
        # 投影到Transformer输入维度
        transformer_input = self.feature_proj(combined_feat)  # 形状: (batch_size, seq_len, d_model)
        # Transformer编码
        encoded_seq = self.transformer_encoder(transformer_input)  # 形状: (batch_size, seq_len, d_model)
        # 取序列最后一步的输出做预测（如果是序列到序列任务，可以取全部步）
        output = self.predictor(encoded_seq[:, -1, :])  # 形状: (batch_size, output_dim)
        return output

输入形状处理说明 #

1. 输入x为(batch_size, seq_len, features)时，先提取单独的时间特征列（或维度），得到(batch_size, seq_len, 1)的张量；
2. 用Time2Vec层编码后得到(batch_size, seq_len, t2v_output_dim)的编码特征；
3. 将编码特征与原输入特征在最后一维拼接，得到适配Transformer输入的特征张量；
4. 通过投影层将特征维度转换为Transformer要求的d_model，即可正常送入Transformer处理。

内容的提问来源于stack exchange，提问作者thunder80