五个时间序列预测的深度学习模型对比总结

class Patch(nn.Module):
def __init__(self, patch_size, input_channel, output_channel, spectral=True):
super().__init__()
self.output_channel = output_channel
self.P = patch_size
self.input_channel = input_channel
self.output_channel = output_channel
self.spectral = spectral
if spectral:
self.emb_layer = nn.Linear(int(patch_size/2+1)*2, output_channel)
else:
self.input_embedding = nn.Conv2d(input_channel, output_channel, kernel_size=(self.P, 1), stride=(self.P, 1))
def forward(self, input):
B, N, C, L = input.shape
if self.spectral:
spec_feat_ = torch.fft.rfft(input.unfold(-1, self.P, self.P), dim=-1)
real = spec_feat_.real
imag = spec_feat_.imag
spec_feat = torch.cat([real, imag], dim=-1).squeeze(2)
output = self.emb_layer(spec_feat).transpose(-1, -2)
else:
input = input.unsqueeze(-1) # B, N, C, L, 1
input = input.reshape(B*N, C, L, 1) # B*N, C, L, 1
output = self.input_embedding(input) # B*N, d, L/P, 1
output = output.squeeze(-1).view(B, N, self.output_channel, -1)
assert output.shape[-1] == L / self.P
return output

class MaskGenerator(nn.Module):
def __init__(self, mask_size, mask_ratio, distribution='uniform', lm=-1):
super().__init__()
self.mask_size = mask_size
self.mask_ratio = mask_ratio
self.sort = True
self.average_patch = lm
self.distribution = distribution
if self.distribution == "geom":
assert lm != -1
assert distribution in ['geom', 'uniform']
def uniform_rand(self):
mask = list(range(int(self.mask_size)))
random.shuffle(mask)
mask_len = int(self.mask_size * self.mask_ratio)
self.masked_tokens = mask[:mask_len]
self.unmasked_tokens = mask[mask_len:]
if self.sort:
self.masked_tokens = sorted(self.masked_tokens)
self.unmasked_tokens = sorted(self.unmasked_tokens)
return self.unmasked_tokens, self.masked_tokens
def geometric_rand(self):
mask = geom_noise_mask_single(self.mask_size, lm=self.average_patch, masking_ratio=self.mask_ratio) # 1: masked, 0:unmasked
self.masked_tokens = np.where(mask)[0].tolist()
self.unmasked_tokens = np.where(~mask)[0].tolist()
# assert len(self.masked_tokens) > len(self.unmasked_tokens)
return self.unmasked_tokens, self.masked_tokens
def forward(self):
if self.distribution == 'geom':
self.unmasked_tokens, self.masked_tokens = self.geometric_rand()
elif self.distribution == 'uniform':
self.unmasked_tokens, self.masked_tokens = self.uniform_rand()
else:
raise Exception("ERROR")
return self.unmasked_tokens, self.masked_tokens

class LearnableTemporalPositionalEncoding(nn.Module):
def __init__(self, d_model, dropout=0.1, max_len: int = 1000):
super().__init__()
self.dropout = nn.Dropout(p=dropout)
self.pe = nn.Parameter(torch.empty(max_len, d_model), requires_grad=True)
nn.init.uniform_(self.pe, -0.02, 0.02)

def forward(self, X, index):
if index is None:
pe = self.pe[:X.size(1), :].unsqueeze(0)
else:
pe = self.pe[index].unsqueeze(0)
X = X + pe
X = self.dropout(X)
return X
class PositionalEncoding(nn.Module):
def __init__(self, hidden_dim, dropout=0.1):
super().__init__()
self.tem_pe = LearnableTemporalPositionalEncoding(hidden_dim, dropout)
def forward(self, input, index=None, abs_idx=None):
B, N, L_P, d = input.shape
# temporal embedding
input = self.tem_pe(input.view(B*N, L_P, d), index=index)
input = input.view(B, N, L_P, d)
# absolute positional embedding
return input

class TransformerLayers(nn.Module):
def __init__(self, hidden_dim, nlayers, num_heads=4, dropout=0.1):
super().__init__()
self.d_model = hidden_dim
encoder_layers = TransformerEncoderLayer(hidden_dim, num_heads, hidden_dim*4, dropout)
self.transformer_encoder = TransformerEncoder(encoder_layers, nlayers)
def forward(self, src):
B, N, L, D = src.shape
src = src * math.sqrt(self.d_model)
src = src.view(B*N, L, D)
src = src.transpose(0, 1)
output = self.transformer_encoder(src, mask=None)
output = output.transpose(0, 1).view(B, N, L, D)
return output

def trunc_normal_(tensor, mean=0., std=1.):
__call_trunc_normal_(tensor, mean=mean, std=std, a=-std, b=std)
def unshuffle(shuffled_tokens):
dic = {}
for k, v, in enumerate(shuffled_tokens):
dic[v] = k
unshuffle_index = []
for i in range(len(shuffled_tokens)):
unshuffle_index.append(dic[i])
return unshuffle_index
class TSFormer(nn.Module):
def __init__(self, patch_size, in_channel, out_channel, dropout, mask_size, mask_ratio, L=6, distribution='uniform', lm=-1, selected_feature=0, mode='Pretrain', spectral=True):
super().__init__()
self.patch_size = patch_size
self.seleted_feature = selected_feature
self.mode = mode
self.spectral = spectral
self.patch = Patch(patch_size, in_channel, out_channel, spectral=spectral)
self.pe = PositionalEncoding(out_channel, dropout=dropout)
self.mask = MaskGenerator(mask_size, mask_ratio, distribution=distribution, lm=lm)
self.encoder = TransformerLayers(out_channel, L)
self.decoder = TransformerLayers(out_channel, 1)
self.encoder_2_decoder = nn.Linear(out_channel, out_channel)
self.mask_token = nn.Parameter(torch.zeros(1, 1, 1, out_channel))
trunc_normal_(self.mask_token, std=.02)
if self.spectral:
self.output_layer = nn.Linear(out_channel, int(patch_size/2+1)*2)
else:
self.output_layer = nn.Linear(out_channel, patch_size)
def _forward_pretrain(self, input):
B, N, C, L = input.shape
# get patches and exec input embedding
patches = self.patch(input) 
patches = patches.transpose(-1, -2) 
# positional embedding
patches = self.pe(patches)

# mask tokens
unmasked_token_index, masked_token_index = self.mask()
encoder_input = patches[:, :, unmasked_token_index, :] 
# encoder
H = self.encoder(encoder_input) 
# encoder to decoder
H = self.encoder_2_decoder(H)
# decoder
# H_unmasked = self.pe(H, index=unmasked_token_index)
H_unmasked = H
H_masked = self.pe(self.mask_token.expand(B, N, len(masked_token_index), H.shape[-1]), index=masked_token_index)
H_full = torch.cat([H_unmasked, H_masked], dim=-2) # # B, N, L/P, d
H = self.decoder(H_full)
# output layer
if self.spectral:
# output = H
spec_feat_H_ = self.output_layer(H)
real = spec_feat_H_[..., :int(self.patch_size/2+1)]
imag = spec_feat_H_[..., int(self.patch_size/2+1):]
spec_feat_H = torch.complex(real, imag)
out_full = torch.fft.irfft(spec_feat_H)
else:
out_full = self.output_layer(H)
# prepare loss
B, N, _, _ = out_full.shape 
out_masked_tokens = out_full[:, :, len(unmasked_token_index):, :]
out_masked_tokens = out_masked_tokens.view(B, N, -1).transpose(1, 2)
label_full = input.permute(0, 3, 1, 2).unfold(1, self.patch_size, self.patch_size)[:, :, :, self.seleted_feature, :].transpose(1, 2) # B, N, L/P, P
label_masked_tokens = label_full[:, :, masked_token_index, :].contiguous()
label_masked_tokens = label_masked_tokens.view(B, N, -1).transpose(1, 2)
# prepare plot
## note that the output_full and label_full are not aligned. The out_full in shuffled
### therefore, unshuffle for plot
unshuffled_index = unshuffle(unmasked_token_index + masked_token_index)
out_full_unshuffled = out_full[:, :, unshuffled_index, :]
plot_args = {}
plot_args['out_full_unshuffled'] = out_full_unshuffled
plot_args['label_full'] = label_full
plot_args['unmasked_token_index'] = unmasked_token_index
plot_args['masked_token_index'] = masked_token_index
return out_masked_tokens, label_masked_tokens, plot_args
def _forward_backend(self, input):
B, N, C, L = input.shape
# get patches and exec input embedding
patches = self.patch(input) 
patches = patches.transpose(-1, -2) 
# positional embedding
patches = self.pe(patches)
encoder_input = patches # no mask when running the backend.
# encoder
H = self.encoder(encoder_input) 
return H
def forward(self, input_data):

if self.mode == 'Pretrain':
return self._forward_pretrain(input_data)
else:
return self._forward_backend(input_data)