mirror of
https://github.com/KimMeen/Time-LLM.git
synced 2025-02-17 17:01:54 +08:00
164 lines
6.3 KiB
Python
164 lines
6.3 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
import matplotlib.pyplot as plt
|
|
import numpy as np
|
|
import math
|
|
from math import sqrt
|
|
import os
|
|
|
|
|
|
class AutoCorrelation(nn.Module):
|
|
"""
|
|
AutoCorrelation Mechanism with the following two phases:
|
|
(1) period-based dependencies discovery
|
|
(2) time delay aggregation
|
|
This block can replace the self-attention family mechanism seamlessly.
|
|
"""
|
|
|
|
def __init__(self, mask_flag=True, factor=1, scale=None, attention_dropout=0.1, output_attention=False):
|
|
super(AutoCorrelation, self).__init__()
|
|
self.factor = factor
|
|
self.scale = scale
|
|
self.mask_flag = mask_flag
|
|
self.output_attention = output_attention
|
|
self.dropout = nn.Dropout(attention_dropout)
|
|
|
|
def time_delay_agg_training(self, values, corr):
|
|
"""
|
|
SpeedUp version of Autocorrelation (a batch-normalization style design)
|
|
This is for the training phase.
|
|
"""
|
|
head = values.shape[1]
|
|
channel = values.shape[2]
|
|
length = values.shape[3]
|
|
# find top k
|
|
top_k = int(self.factor * math.log(length))
|
|
mean_value = torch.mean(torch.mean(corr, dim=1), dim=1)
|
|
index = torch.topk(torch.mean(mean_value, dim=0), top_k, dim=-1)[1]
|
|
weights = torch.stack([mean_value[:, index[i]] for i in range(top_k)], dim=-1)
|
|
# update corr
|
|
tmp_corr = torch.softmax(weights, dim=-1)
|
|
# aggregation
|
|
tmp_values = values
|
|
delays_agg = torch.zeros_like(values).float()
|
|
for i in range(top_k):
|
|
pattern = torch.roll(tmp_values, -int(index[i]), -1)
|
|
delays_agg = delays_agg + pattern * \
|
|
(tmp_corr[:, i].unsqueeze(1).unsqueeze(1).unsqueeze(1).repeat(1, head, channel, length))
|
|
return delays_agg
|
|
|
|
def time_delay_agg_inference(self, values, corr):
|
|
"""
|
|
SpeedUp version of Autocorrelation (a batch-normalization style design)
|
|
This is for the inference phase.
|
|
"""
|
|
batch = values.shape[0]
|
|
head = values.shape[1]
|
|
channel = values.shape[2]
|
|
length = values.shape[3]
|
|
# index init
|
|
init_index = torch.arange(length).unsqueeze(0).unsqueeze(0).unsqueeze(0).repeat(batch, head, channel, 1).cuda()
|
|
# find top k
|
|
top_k = int(self.factor * math.log(length))
|
|
mean_value = torch.mean(torch.mean(corr, dim=1), dim=1)
|
|
weights, delay = torch.topk(mean_value, top_k, dim=-1)
|
|
# update corr
|
|
tmp_corr = torch.softmax(weights, dim=-1)
|
|
# aggregation
|
|
tmp_values = values.repeat(1, 1, 1, 2)
|
|
delays_agg = torch.zeros_like(values).float()
|
|
for i in range(top_k):
|
|
tmp_delay = init_index + delay[:, i].unsqueeze(1).unsqueeze(1).unsqueeze(1).repeat(1, head, channel, length)
|
|
pattern = torch.gather(tmp_values, dim=-1, index=tmp_delay)
|
|
delays_agg = delays_agg + pattern * \
|
|
(tmp_corr[:, i].unsqueeze(1).unsqueeze(1).unsqueeze(1).repeat(1, head, channel, length))
|
|
return delays_agg
|
|
|
|
def time_delay_agg_full(self, values, corr):
|
|
"""
|
|
Standard version of Autocorrelation
|
|
"""
|
|
batch = values.shape[0]
|
|
head = values.shape[1]
|
|
channel = values.shape[2]
|
|
length = values.shape[3]
|
|
# index init
|
|
init_index = torch.arange(length).unsqueeze(0).unsqueeze(0).unsqueeze(0).repeat(batch, head, channel, 1).cuda()
|
|
# find top k
|
|
top_k = int(self.factor * math.log(length))
|
|
weights, delay = torch.topk(corr, top_k, dim=-1)
|
|
# update corr
|
|
tmp_corr = torch.softmax(weights, dim=-1)
|
|
# aggregation
|
|
tmp_values = values.repeat(1, 1, 1, 2)
|
|
delays_agg = torch.zeros_like(values).float()
|
|
for i in range(top_k):
|
|
tmp_delay = init_index + delay[..., i].unsqueeze(-1)
|
|
pattern = torch.gather(tmp_values, dim=-1, index=tmp_delay)
|
|
delays_agg = delays_agg + pattern * (tmp_corr[..., i].unsqueeze(-1))
|
|
return delays_agg
|
|
|
|
def forward(self, queries, keys, values, attn_mask):
|
|
B, L, H, E = queries.shape
|
|
_, S, _, D = values.shape
|
|
if L > S:
|
|
zeros = torch.zeros_like(queries[:, :(L - S), :]).float()
|
|
values = torch.cat([values, zeros], dim=1)
|
|
keys = torch.cat([keys, zeros], dim=1)
|
|
else:
|
|
values = values[:, :L, :, :]
|
|
keys = keys[:, :L, :, :]
|
|
|
|
# period-based dependencies
|
|
q_fft = torch.fft.rfft(queries.permute(0, 2, 3, 1).contiguous(), dim=-1)
|
|
k_fft = torch.fft.rfft(keys.permute(0, 2, 3, 1).contiguous(), dim=-1)
|
|
res = q_fft * torch.conj(k_fft)
|
|
corr = torch.fft.irfft(res, dim=-1)
|
|
|
|
# time delay agg
|
|
if self.training:
|
|
V = self.time_delay_agg_training(values.permute(0, 2, 3, 1).contiguous(), corr).permute(0, 3, 1, 2)
|
|
else:
|
|
V = self.time_delay_agg_inference(values.permute(0, 2, 3, 1).contiguous(), corr).permute(0, 3, 1, 2)
|
|
|
|
if self.output_attention:
|
|
return (V.contiguous(), corr.permute(0, 3, 1, 2))
|
|
else:
|
|
return (V.contiguous(), None)
|
|
|
|
|
|
class AutoCorrelationLayer(nn.Module):
|
|
def __init__(self, correlation, d_model, n_heads, d_keys=None,
|
|
d_values=None):
|
|
super(AutoCorrelationLayer, self).__init__()
|
|
|
|
d_keys = d_keys or (d_model // n_heads)
|
|
d_values = d_values or (d_model // n_heads)
|
|
|
|
self.inner_correlation = correlation
|
|
self.query_projection = nn.Linear(d_model, d_keys * n_heads)
|
|
self.key_projection = nn.Linear(d_model, d_keys * n_heads)
|
|
self.value_projection = nn.Linear(d_model, d_values * n_heads)
|
|
self.out_projection = nn.Linear(d_values * n_heads, d_model)
|
|
self.n_heads = n_heads
|
|
|
|
def forward(self, queries, keys, values, attn_mask):
|
|
B, L, _ = queries.shape
|
|
_, S, _ = keys.shape
|
|
H = self.n_heads
|
|
|
|
queries = self.query_projection(queries).view(B, L, H, -1)
|
|
keys = self.key_projection(keys).view(B, S, H, -1)
|
|
values = self.value_projection(values).view(B, S, H, -1)
|
|
|
|
out, attn = self.inner_correlation(
|
|
queries,
|
|
keys,
|
|
values,
|
|
attn_mask
|
|
)
|
|
out = out.view(B, L, -1)
|
|
|
|
return self.out_projection(out), attn
|