import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.utils import weight_norm
import math
#定义了一个位置嵌入类,它能提供序列中每个元素的位置信息。
class PositionalEmbedding(nn.Module):
def __init__(self, d_model, max_len=5000):
# super(PositionalEmbedding, self).__init__() 在子类中调用父类的初始化方法。
# pe = torch.zeros(max_len, d_model).float() 创建一个max_len行d_model列的浮点数零张量,用于存储位置嵌入。
# pe.require_grad = False 设置这个张量不需要计算梯度,因为位置嵌入是固定的。
# position = torch.arange(0, max_len).float().unsqueeze(1) 创建一个0到max_len的连续值,并增加一个维度。
# div_term = ... 计算位置编码的分频项。
# pe[:, 0::2] = torch.sin(position * div_term) 使用正弦函数给偶数位置编码。
# pe[:, 1::2] = torch.cos(position * div_term) 使用余弦函数给奇数位置编码。
# pe = pe.unsqueeze(0) 增加一维,为批量维度做准备。
# self.register_buffer('pe', pe) 将位置嵌入张量注册为模型的一个缓冲区,它不会被认为是一个模型参数。
super(PositionalEmbedding, self).__init__()
# Compute the positional encodings once in log space.
pe = torch.zeros(max_len, d_model).float()
pe.require_grad = False
position = torch.arange(0, max_len).float().unsqueeze(1)
div_term = (torch.arange(0, d_model, 2).float()
* -(math.log(10000.0) / d_model)).exp()
pe[:, 0::2] = torch.sin(position * div_term)
pe[:, 1::2] = torch.cos(position * div_term)
pe = pe.unsqueeze(0)
self.register_buffer('pe', pe)
def forward(self, x):
# return self.pe[:, :x.size(1)] 在前向传播时返回位置嵌入的一个子集,其长度与输入x的长度一致。
return self.pe[:, :x.size(1)]
class TokenEmbedding(nn.Module):
def __init__(self, c_in, d_model):
# super(TokenEmbedding, self).__init__() 调用父类的初始化方法。
# padding = ... 根据PyTorch版本来决定卷积的填充值。
# self.tokenConv = nn.Conv1d(...) 定义一个一维卷积层,用于从输入序列中学习表示。
super(TokenEmbedding, self).__init__()
padding = 1 if torch.__version__ >= '1.5.0' else 2
self.tokenConv = nn.Conv1d(in_channels=c_in, out_channels=d_model,
kernel_size=3, padding=padding, padding_mode='circular', bias=False)
for m in self.modules():
if isinstance(m, nn.Conv1d):
# nn.init.kaiming_normal_(...) 使用Kaiming初始化方法初始化卷积层的权重。
nn.init.kaiming_normal_(
m.weight, mode='fan_in', nonlinearity='leaky_relu')
def forward(self, x):
x = self.tokenConv(x.permute(0, 2, 1)).transpose(1, 2) #在前向传播中,对输入x进行维度交换,应用卷积,然后再次交换维度,使其回到原始的形状。
return x
class FixedEmbedding(nn.Module):
def __init__(self, c_in, d_model):
# super(FixedEmbedding, self).__init__() 调用父类的初始化方法。
# w = torch.zeros(c_in, d_model).float() 创建一个用于存储嵌入的零张量。
# w.require_grad = False 设置不需要计算梯度。
# self.emb = nn.Embedding(c_in, d_model) 创建一个嵌入层。
# self.emb.weight = nn.Parameter(w, requires_grad=False) 将预计算好的权重作为嵌入层的权重。
super(FixedEmbedding, self).__init__()
w = torch.zeros(c_in, d_model).float()
w.require_grad = False
position = torch.arange(0, c_in).float().unsqueeze(1)
div_term = (torch.arange(0, d_model, 2).float()
* -(math.log(10000.0) / d_model)).exp()
w[:, 0::2] = torch.sin(position * div_term)
w[:, 1::2] = torch.cos(position * div_term)
self.emb = nn.Embedding(c_in, d_model)
self.emb.weight = nn.Parameter(w, requires_grad=False)
def forward(self, x):
# return self.emb(x).detach() 返回嵌入的结果,并确保不会计算梯度。
return self.emb(x).detach()
#定义了一个时间嵌入类,用于处理与时间相关的输入特征
class TemporalEmbedding(nn.Module):
def __init__(self, d_model, embed_type='fixed', freq='h'):
#super(TemporalEmbedding, self).__init__() 调用父类的初始化方法。
#minute_size、hour_size 等定义了时间相关特征的可能值的大小。
#Embed = FixedEmbedding if embed_type == 'fixed' else nn.Embedding 根据嵌入类型选择使用固定嵌入还是可学习的嵌入。
#self.minute_embed 等创建不同时间维度的嵌入层。
super(TemporalEmbedding, self).__init__()
minute_size = 4
hour_size = 24
weekday_size = 7
day_size = 32
month_size = 13
Embed = FixedEmbedding if embed_type == 'fixed' else nn.Embedding
if freq == 't':
self.minute_embed = Embed(minute_size, d_model)
self.hour_embed = Embed(hour_size, d_model)
self.weekday_embed = Embed(weekday_size, d_model)
self.day_embed = Embed(day_size, d_model)
self.month_embed = Embed(month_size, d_model)
def forward(self, x):
x = x.long()
minute_x = self.minute_embed(x[:, :, 4]) if hasattr(
self, 'minute_embed') else 0.
hour_x = self.hour_embed(x[:, :, 3])
weekday_x = self.weekday_embed(x[:, :, 2])
day_x = self.day_embed(x[:, :, 1])
month_x = self.month_embed(x[:, :, 0])
#return hour_x + weekday_x + day_x + month_x + minute_x 在前向传播时返回所有时间嵌入的总和。
return hour_x + weekday_x + day_x + month_x + minute_x
#定义了一个时间特征嵌入类,仅使用一个线性层来嵌入时间特征。
class TimeFeatureEmbedding(nn.Module):
def __init__(self, d_model, embed_type='timeF', freq='h'):
# super(TimeFeatureEmbedding, self).__init__() 调用父类的初始化方法。
# freq_map = ... 定义了不同频率对应的输入维度。
# self.embed = nn.Linear(d_inp, d_model, bias=False) 创建一个线性层,用于将时间特征投影到一个d_model维的空间。
super(TimeFeatureEmbedding, self).__init__()
freq_map = {'h': 4, 't': 5, 's': 6,
'm': 1, 'a': 1, 'w': 2, 'd': 3, 'b': 3}
d_inp = freq_map[freq]
self.embed = nn.Linear(d_inp, d_model, bias=False)
def forward(self, x):
# return self.embed(x) 返回线性层的输出。
return self.embed(x)
class DataEmbedding(nn.Module):
def __init__(self, c_in, d_model, embed_type='fixed', freq='h', dropout=0.1):
# super(DataEmbedding, self).__init__() 调用父类的初始化方法。
# self.value_embedding = ... 创建一个用于值的嵌入层。
# self.position_embedding = ... 创建一个位置嵌入实例。
# self.temporal_embedding = ... 创建一个时间嵌入实例。
# self.dropout = nn.Dropout(p=dropout) 定义一个dropout层,用于防止过拟合。
super(DataEmbedding, self).__init__()
self.value_embedding = TokenEmbedding(c_in=c_in, d_model=d_model)
self.position_embedding = PositionalEmbedding(d_model=d_model)
self.temporal_embedding = TemporalEmbedding(d_model=d_model, embed_type=embed_type,
freq=freq) if embed_type != 'timeF' else TimeFeatureEmbedding(
d_model=d_model, embed_type=embed_type, freq=freq)
self.dropout = nn.Dropout(p=dropout)
def forward(self, x, x_mark):
# if x_mark is None: ... 根据是否有时间标记来决定如何计算嵌入。
if x_mark is None:
x = self.value_embedding(x) + self.position_embedding(x)
else:
x = self.value_embedding(
x) + self.temporal_embedding(x_mark) + self.position_embedding(x)
return self.dropout(x)
class DataEmbedding_inverted(nn.Module):
def __init__(self, c_in, d_model, embed_type='fixed', freq='h', dropout=0.1):
super(DataEmbedding_inverted, self).__init__()
self.value_embedding = nn.Linear(c_in, d_model)
self.dropout = nn.Dropout(p=dropout)
def forward(self, x, x_mark):
x = x.permute(0, 2, 1)
# x: [Batch Variate Time]
if x_mark is None:
x = self.value_embedding(x)
else:
x = self.value_embedding(torch.cat([x, x_mark.permute(0, 2, 1)], 1))
# x: [Batch Variate d_model]
return self.dropout(x)
#定义了一个数据嵌入类,它使用线性层而不是卷积层来嵌入输入数据。
class DataEmbedding_wo_pos(nn.Module):
def __init__(self, c_in, d_model, embed_type='fixed', freq='h', dropout=0.1):
# super(DataEmbedding_inverted, self).__init__() 调用父类的初始化方法。
# self.value_embedding = nn.Linear(c_in, d_model) 创建一个线性层用于值嵌入。
# x = x.permute(0, 2, 1) 在前向传播中,交换维度以适应线性层的输入格式。
super(DataEmbedding_wo_pos, self).__init__()
self.value_embedding = TokenEmbedding(c_in=c_in, d_model=d_model)
self.position_embedding = PositionalEmbedding(d_model=d_model)
self.temporal_embedding = TemporalEmbedding(d_model=d_model, embed_type=embed_type,
freq=freq) if embed_type != 'timeF' else TimeFeatureEmbedding(
d_model=d_model, embed_type=embed_type, freq=freq)
self.dropout = nn.Dropout(p=dropout)
def forward(self, x, x_mark):
# if x_mark is None: ... 根据是否有时间标记来决定如何组合嵌入。
if x_mark is None:
x = self.value_embedding(x)
else:
x = self.value_embedding(x) + self.temporal_embedding(x_mark)
return self.dropout(x)
#定义了一个补丁嵌入类,用于处理将数据切分成补丁并嵌入。
class PatchEmbedding(nn.Module):
def __init__(self, d_model, patch_len, stride, padding, dropout):
# super(PatchEmbedding, self).__init__() 调用父类的初始化方法。
# self.patch_len、self.stride、self.padding_patch_layer 定义了补丁的长度、步长和填充层。
# self.value_embedding = ... 创建一个线性层用于补丁的嵌入。
# self.position_embedding = ... 创建一个位置嵌入实例。
super(PatchEmbedding, self).__init__()
# Patching
self.patch_len = patch_len
self.stride = stride
self.padding_patch_layer = nn.ReplicationPad1d((0, padding))
# Backbone, Input encoding: projection of feature vectors onto a d-dim vector space
self.value_embedding = nn.Linear(patch_len, d_model, bias=False)
# Positional embedding
self.position_embedding = PositionalEmbedding(d_model)
# Residual dropout
self.dropout = nn.Dropout(dropout)
def forward(self, x):
# n_vars = x.shape[1] 获取变量数量。
# x = self.padding_patch_layer(x) 应用填充。
# x = x.unfold(dimension=-1, ...) 将数据展开成补丁。
# x = torch.reshape(x, ...) 重塑张量以适应线性层。
# x = self.value_embedding(x) + self.position_embedding(x) 应用值嵌入和位置嵌入。
# return self.dropout(x), n_vars 返回dropout后的嵌入结果和变量数量。
# do patching
n_vars = x.shape[1]
x = self.padding_patch_layer(x)
x = x.unfold(dimension=-1, size=self.patch_len, step=self.stride)
x = torch.reshape(x, (x.shape[0] * x.shape[1], x.shape[2], x.shape[3]))
# Input encoding
x = self.value_embedding(x) + self.position_embedding(x)
return self.dropout(x), n_vars