Transformer 天气数据进行时序预测

Transformer原理可以参见我转载的这篇文章

手动求解 Transformer：分步数学示例 (转载)-CSDN博客

数据集选择

如果具有采集好的csv 数据文件，文件列需要包含足备数据提供给模型训练使用

下面是我搜集的时序数据集

链接：https://pan.baidu.com/s/1KtLeqsTu_1a5AhBIAq3lQA?pwd=frf5

提取码：frf5

我选择使用天气的数据集做试验

这组数据包含了气象学中常用的一些参数。以下是每一列的含义：

date: 日期，表示数据记录的时间。
p (mbar): 气压，以毫巴（mbar）为单位。
T (degC): 温度，以摄氏度（°C）为单位。
Tpot (K): 潜在温度，以开尔文（K）为单位。
Tdew (degC): 露点温度，以摄氏度（°C）为单位。
rh (%): 相对湿度，以百分比（%）表示。
VPmax (mbar): 最大蒸气压，以毫巴（mbar）为单位。
VPact (mbar): 实际蒸气压，以毫巴（mbar）为单位。
VPdef (mbar): 蒸气压缺乏，以毫巴（mbar）为单位。
sh (g/kg): 比湿，以克每千克（g/kg）为单位。
H2OC (mmol/mol): 水汽含量，以每摩尔气体中毫摩尔（mmol/mol）表示。
rho (g/m³): 密度，以克每立方米（g/m³）为单位。
wv (m/s): 风速，以米每秒（m/s）为单位。
max. wv (m/s): 最大风速，以米每秒（m/s）为单位。
wd (deg): 风向，以度（°）表示。
rain (mm): 降雨量，以毫米（mm）为单位。
raining (s): 下雨的持续时间，以秒（s）为单位。
SWDR (W/m²): 太阳辐射强度，以瓦特每平方米（W/m²）为单位。
PAR (μmol/m²/s): 光合有效辐射，以微摩尔每平方米每秒（μmol/m²/s）为单位。
max. PAR (μmol/m²/s): 最大光合有效辐射，以微摩尔每平方米每秒（μmol/m²/s）为单位。
Tlog (degC): 记录器温度，以摄氏度（°C）为单位。

特征列还是太多了，利用皮尔逊相关系数矩阵来筛选使用的特征，具体可以参见我的这篇笔记CSDN

以下是皮尔逊相关系数的一些要点：

相关系数值范围:
- 1: 完全正相关，表示两个变量之间有完全的线性关系，且一方的增加会导致另一方的增加。
- -1: 完全负相关，表示两个变量之间有完全的线性关系，且一方的增加会导致另一方的减少。
- 0: 无相关性，表示两个变量之间没有线性关系。
靠近1或-1的值:
- 值靠近1 或**-1**表示强相关性。具体来说：
  - 靠近1的值表示强正相关，意味着一个特征可以很好地预测另一个特征。
  - 靠近**-1**的值表示强负相关，意味着一个特征的增加与另一个特征的减少有很强的关系。
预测能力:
- 单独的相关系数不能完全代表预测能力，但高绝对值的相关系数通常意味着可能的线性关系，因此在某种程度上可以提示特征之间的预测能力。

从热力图上看到强预测能力的特征选择

rho (g/m**3)

sh (g/kg)

T (degC)

Tdew (degC)

Tlog (degC)

Tpot (K)

VPact (mbar)

VPmax (mbar)

使用pandas把相应特征提取到pkl文件里

# -*- coding: utf-8 -*-
import pandas as pd
import numpy as np

df=pd.read_excel('./data/weather.xls', engine='xlrd')

'''
rho (g/m**3) 密度
sh (g/kg) 比湿
T (degC) 温度 
Tdew (degC) 露点温度
Tlog (degC) 记录器温度
Tpot (K) 潜在温度
VPact (mbar) 实际蒸气压
VPmax (mbar) 最大蒸气压
'''
data=pd.concat([np.floor((df['rho (g/m**3)'])),
            np.floor((df['sh (g/kg)'])),
            np.floor((df['T (degC)'])),
            np.floor((df['Tdew (degC)'])),
            np.floor((df['Tlog (degC)'])),
            np.floor((df['Tpot (K)'])),
            np.floor((df['VPact (mbar)'])),
            np.floor((df['VPmax (mbar)']))], axis=1)

data.to_pickle('./data/weather.pkl')

将数据拆分为训练集和验证集

import pickle
from torch.utils.data import Dataset
import numpy as np
import torch

with open('./data/weather.pkl','rb') as f:
    data = pickle.load(f)
    data_train_raw, data_test_raw = [], []
    # 数据种包含 52697 条数据， 按照20 step 做一次预测，即每三个小的历史数据实现一次时序预测
    step = 20
    i = 0
    while i < 26000:
        data_train_raw.append(data[i:i + 11])
        i += step
    # 验证集
    i = 26000
    while i < 52000:
        data_test_raw.append(data[i:i + 11])
        i += step
    pass

class WeatherTrajData(Dataset):

    def __init__(self, data):
        self.data = torch.from_numpy(np.array(data)).to(torch.float32)

    def __getitem__(self, index):
        return self.data[index, :, :]

    def __len__(self):
        return self.data.size(0)
        # return data.size(0)

下面要贴点正确的废话，Transformer为什么可以做时序预测

利用Transformer 模型，关注attetion注意力技术，它使得模型可以根据需要，关注到输入序列的相关部分。Attention不止用在NLP领域。换一个更一般点的例子，来说明Attention的一般模式。

首先计算q和k的相似度s1 - sn，并归一化到a1 - an，a反映了孩子对玩具的喜好程度 (权重)。接下来a对特征v进行加权求和 (家长角度考虑)，得到向量h。最后家长是否购买玩具是由向量h决定的。

上述过程就是Attention的标准操作流程。Attention模型三要素是g、K、V。K、V矩阵分别对应向量序列ki到hn和v到n。由于中间涉及到加权求和，所以这两个序列长度一致，而且元素都是对应的。即k;对应vi。但是k和v分别表示两类特征，所以向量长度可以不一致。

为了运算方便，可以将Attention操作计算为: h = Attention (g、K、V)。g也可以是一个向量序列Q(对应机器翻译中输入多个单词)，此时输出也是一个向量序列H。Attention通用标准公式为: $$H=Attention(Q,K,V)= beginbmatrix) Attention(g (1),K,V) ....Attention(q (m),K,V) endfbmatrix;$$

这里，Q、飞、V、H均为矩阵 (向量序列) 。其中， H和Q序列长度一致，各行一一对应(一个输入对应一个输出)，K和V序列长度一致，各行一一对应。

Attention模型计算相似度，除了直接求内积<q,k>，还有很多其它形式

s = ATTanh (gW +kU)

多层感知机，A、W、U都是待学习参数。这种方法不仅避开了求内积时 g 和 k 的向量长度必须致的限制,还可以进行不同空间的向量匹配。例如,在进行图文匹配时,q 和 k 分别来自文字空间和图像空间,可以先分别通过、U将它们转换至同一空间,再求相似度

上面式子中，w、U是矩阵参数，相乘后可以使g和k的维度一致。比如机器翻译中，中文一词多义情况比较多，中文向量q维度可以设置长一点，k的维度可以设置短一点。gw +kU是同长度向量相加，结果还是一个向量，再经过列向量AT相乘，得到一个标量，即attention score数值.

（1）attention的mask是可选的，在时序任务里是需要的，因为时序是有因果性的。

（2）Q*K才是attention，它的维度是Batch*heads*lens*lens，也就意味着是有明确物理意义的，只和最初的序列长度有关系，可以可视化出来做分析用。

（3）ScaledDotProduct算完后，最后一维是由dv决定的，想回到q的维度，还需要做linear的操作。

（4）contiguous & view：代码里的操作实际就是把dv数据按照heads concate到了一起，具体为何需要使用contiguous的原因PyTorch中的contiguous 。

（5）residual：

残差连接通常是在每个子层（如自注意力层和前馈神经网络层）之后进行的。具体来说，输入 xxx 会先经过某个神经网络层（如自注意力层），然后与输入 xxx 相加，形成输出：

这里简单给出整体数据流和主要的模型结构

利用chatgpt我们来攒一下transformer的模型代码

编码器和解码器

class Encoder(nn.Module):
    ''' A encoder model with self attention mechanism. '''

    def __init__(
            self, n_src_vocab, d_word_vec, n_layers, n_head, d_k, d_v,
            d_model, d_inner, dropout=0.1, n_position=200, scale_emb=False):

        super().__init__()

        self.src_word_emb = nn.Embedding(n_src_vocab, d_word_vec)
        self.position_enc = PositionalEncoding(d_word_vec, n_position=n_position)
        self.dropout = nn.Dropout(p=dropout)
        self.layer_stack = nn.ModuleList([
            EncoderLayer(d_model, d_inner, n_head, d_k, d_v, dropout=dropout)
            for _ in range(n_layers)])
        self.layer_norm = nn.LayerNorm(d_model, eps=1e-6)
        self.scale_emb = scale_emb
        self.d_model = d_model

    def forward(self, src_seq, return_attns=False):

        enc_slf_attn_list = []

        # -- Forward
        enc_output = self.src_word_emb(src_seq.long())
        if self.scale_emb:
            enc_output *= self.d_model ** 0.5
        enc_output = self.dropout(self.position_enc(enc_output))
        enc_output = self.layer_norm(enc_output)

        for enc_layer in self.layer_stack:
            enc_output, enc_slf_attn = enc_layer(enc_output)
            enc_slf_attn_list += [enc_slf_attn] if return_attns else []

        if return_attns:
            return enc_output, enc_slf_attn_list
        return enc_output


class Decoder(nn.Module):
    ''' A decoder model with self attention mechanism. '''

    def __init__(
            self, n_trg_vocab, d_word_vec, n_layers, n_head, d_k, d_v,
            d_model, d_inner, pad_idx, n_position=200, dropout=0.1, scale_emb=False):

        super().__init__()

        self.trg_word_emb = nn.Embedding(n_trg_vocab, d_word_vec, padding_idx=pad_idx)
        self.position_enc = PositionalEncoding(d_word_vec, n_position=n_position)
        self.dropout = nn.Dropout(p=dropout)
        self.layer_stack = nn.ModuleList([
            DecoderLayer(d_model, d_inner, n_head, d_k, d_v, dropout=dropout)
            for _ in range(n_layers)])
        self.layer_norm = nn.LayerNorm(d_model, eps=1e-6)
        self.scale_emb = scale_emb
        self.d_model = d_model

    def forward(self, trg_seq, trg_mask, enc_output, return_attns=False):

        dec_slf_attn_list, dec_enc_attn_list = [], []
        dec_output = self.trg_word_emb(trg_seq.long())
        if self.scale_emb:
            dec_output *= self.d_model ** 0.5
        dec_output = self.dropout(self.position_enc(dec_output))
        dec_output = self.layer_norm(dec_output)

        for dec_layer in self.layer_stack:
            dec_output, dec_slf_attn, dec_enc_attn = dec_layer(
                dec_output, enc_output)
            dec_slf_attn_list += [dec_slf_attn] if return_attns else []
            dec_enc_attn_list += [dec_enc_attn] if return_attns else []

        if return_attns:
            return dec_output, dec_slf_attn_list, dec_enc_attn_list
        return dec_output

位置编码

class PositionalEncoding(nn.Module):

    def __init__(self, d_hid, n_position=200):
        super(PositionalEncoding, self).__init__()

        # Not a parameter
        '''
        存储非学习参数:通常情况下,我们会将神经网络中的可学习参数(权重、偏移等)注册为 nn.Parameter，
        而一些非学习参数(如均值、方差等)则可以使用 register_buffer 进行注册。这样可以将这些非学习参数与模型本身进行绑定,方便后续的保存和加载。
        减少显存占用:相比于 nn.Parameter，register_buffer 注册的参数不会作为模型的可学习参数进行反向传播更新。
        这样可以减少显存占用,特别是在处理大型模型时非常有用。
        方便模型保存和加载:当我们保存和加载模型时,使用 register_buffer 注册的参数会自动被包含在模型的状态字典中,避免了手动处理这些参数的麻烦。
        增强模型可解释性:通过 register_buffer 注册一些中间计算结果或者辅助性参数,可以帮助理解模型的内部工作机制,增强模型的可解释性。
        '''
        self.register_buffer('pos_table', self._get_sinusoid_encoding_table(n_position, d_hid))

    def _get_sinusoid_encoding_table(self, n_position, d_hid):
        ''' Sinusoid position encoding table '''
        '''
        作用是生成一个用于位置编码的正弦波表。位置编码是自然语言处理中常用的技术之一,
        它的目的是为输入序列中的每个词添加一个位置信息,从而帮助模型理解词语在序列中的相对位置关系。
        具体来说:
            n_position 表示序列的最大长度,d_hid 表示每个词的隐藏层维度。
            该方法首先定义了一个 get_position_angle_vec 辅助函数,用于计算每个位置对应的正弦波值。
            然后使用 NumPy 生成一个 n_position x d_hid 大小的正弦波编码表 sinusoid_table。其中奇数列存储 sine 值,偶数列存储 cosine 值。
            最后将 sinusoid_table 转换为 PyTorch 张量并增加一个 batch 维度,得到最终的位置编码表。
            这个位置编码表可以作为模型输入的一部分,与词嵌入向量进行拼接,从而为模型提供额外的位置信息。
            相比于学习可训练的位置编码,这种基于正弦波的预定义位置编码在某些情况下可以取得不错的效果,同时计算代价也较低。
        '''
        def get_position_angle_vec(position):
            return [position / np.power(10000, 2 * (hid_j // 2) / d_hid) for hid_j in range(d_hid)]

        sinusoid_table = np.array([get_position_angle_vec(pos_i) for pos_i in range(n_position)])
        sinusoid_table[:, 0::2] = np.sin(sinusoid_table[:, 0::2])  # dim 2i
        sinusoid_table[:, 1::2] = np.cos(sinusoid_table[:, 1::2])  # dim 2i+1

        return torch.FloatTensor(sinusoid_table).unsqueeze(0)

    def forward(self, x):
        return x + self.pos_table[:, :x.size(1)].clone().detach()

多头注意力

class MultiHeadAttention(nn.Module):
    ''' Multi-Head Attention module '''

    def __init__(self, n_head, d_model, d_k, d_v, dropout=0.1):
        super().__init__()

        self.n_head = n_head
        self.d_k = d_k
        self.d_v = d_v

        self.w_qs = nn.Linear(d_model, n_head * d_k, bias=False)
        self.w_ks = nn.Linear(d_model, n_head * d_k, bias=False)
        self.w_vs = nn.Linear(d_model, n_head * d_v, bias=False)
        self.fc = nn.Linear(n_head * d_v, d_model, bias=False)

        self.attention = ScaledDotProductAttention(temperature=d_k ** 0.5)

        self.dropout = nn.Dropout(dropout)
        self.layer_norm = nn.LayerNorm(d_model, eps=1e-6)


    def forward(self, q, k, v, mask=None):

        d_k, d_v, n_head = self.d_k, self.d_v, self.n_head
        sz_b, len_q, len_k, len_v = q.size(0), q.size(1), k.size(1), v.size(1)

        residual = q
        # Pass through the pre-attention projection: b x lq x (n*dv)
        # Separate different heads: b x lq x n x dv
        q = self.w_qs(q).view(sz_b, len_q, n_head, d_k)
        k = self.w_ks(k).view(sz_b, len_k, n_head, d_k)
        v = self.w_vs(v).view(sz_b, len_v, n_head, d_v)

        # Transpose for attention dot product: b x n x lq x dv
        q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)

        if mask is not None:
            mask = mask.unsqueeze(1)   # For head axis broadcasting.

        q, attn = self.attention(q, k, v, mask=mask)

        # Transpose to move the head dimension back: b x lq x n x dv
        # Combine the last two dimensions to concatenate all the heads together: b x lq x (n*dv)
        q = q.transpose(1, 2).contiguous().view(sz_b, len_q, -1)
        q = self.dropout(self.fc(q))
        q += residual

        q = self.layer_norm(q)

        return q, attn

Transformer 模型

class Transformer(nn.Module):
    ''' A sequence to sequence model with attention mechanism. '''

    '''
    n_src_vocab：源语言词汇表大小。表示源语言中不同词的个数。
    n_trg_vocab：目标语言词汇表大小。表示目标语言中不同词的个数。
    d_word_vec：词嵌入维度。表示将每个词映射到固定维度的向量表示的大小。在Transformer模型中，源语言和目标语言的词嵌入维度可以是不同的。
    d_model：模型的隐藏层维度。表示Transformer模型中每个位置的隐藏层表示的大小。它也是词嵌入维度和注意力机制中的查询、键、值向量的维度。
    d_inner：前馈神经网络内部层的维度。表示Transformer模型中前馈神经网络内部隐藏层的大小。
    n_layers：模型的层数。表示Transformer模型中的编码器和解码器堆叠的层数。
    n_head：注意力头的数量。表示Transformer模型中多头自注意力机制的头数，用于捕捉不同的语义信息。
    d_k：注意力机制中的查询（Q）和键（K）向量的维度。
    d_v：注意力机制中的值（V）向量的维度。
    dropout：Dropout概率。表示在模型训练过程中应用的Dropout概率，用于减少过拟合。
    n_position：位置编码的最大长度。表示Transformer模型中位置编码的最大长度，用于表示输入序列中不同位置的相对位置信息。
    '''

    def __init__(
            self,
            n_src_vocab,
            n_trg_vocab,
            d_word_vec=512,
            d_model=512,
            d_inner=2048,
            n_layers=6,
            n_head=8,
            d_k=64,
            d_v=64,
            dropout=0.1,
            n_position=200):

        super().__init__()
        self.d_model = d_model
        self.fpn = FPN()
        self.encoder = Encoder(
            n_src_vocab=n_src_vocab, n_position=n_position,
            d_word_vec=d_word_vec, d_model=d_model, d_inner=d_inner,
            n_layers=n_layers, n_head=n_head, d_k=d_k, d_v=d_v,
            dropout=dropout)

        self.decoder = Decoder(
            n_trg_vocab=n_trg_vocab, n_position=n_position,
            d_word_vec=d_word_vec, d_model=d_model, d_inner=d_inner,
            n_layers=n_layers, pad_idx=1, n_head=n_head, d_k=d_k, d_v=d_v,
            dropout=dropout)

        self.fdn = FDN()

        self.trg_word_prj = nn.Linear(d_model, n_trg_vocab, bias=False)

        '''
        nn.init.xavier_uniform_ 是 PyTorch 中的一个参数初始化方法，用于初始化神经网络模型的参数。
        具体来说，它使用 Xavier（也称为Glorot）初始化方法，对权重矩阵进行均匀分布的初始化。
        Xavier初始化方法旨在使得输入和输出的方差保持一致，从而避免在深度神经网络中出现梯度消失或梯度爆炸的问题。
        这种初始化方法假设输入和输出的方差相等，并且通过适当的缩放因子将权重初始化在一个合适的范围内。
        nn.init.xavier_uniform_ 方法通过从均匀分布中采样权重值，并根据输入和输出维度进行缩放，将初始化的权重赋值给模型的参数。
        这样做可以确保模型的初始权重适合于有效的前向传播和反向传播。
        '''
        for p in self.parameters():
            if p.dim() > 1:
                nn.init.xavier_uniform_(p)

    def forward(self, input_data, device):
        input_data = self.fpn(input_data).squeeze(-1)
        # 除了最后一行以外的切片
        src_seq = input_data[:, :-1].to(device).to(torch.float32)
        # 除了第一行以外的切片
        trg_seq = input_data[:, 1:].to(device).to(torch.float32)
        trg_mask = get_subsequent_mask(trg_seq)
        enc_output = self.encoder(src_seq)
        dec_output = self.decoder(trg_seq, trg_mask, enc_output)
        trajectory_logit = self.trg_word_prj(dec_output)
        trajectory_logit = self.fdn(trajectory_logit)

        return trajectory_logit

在transformer 模型的输入和输出需要加上两层神经网络，实现训练时把多个特征融合为一个，预测时候把transformer的目标维度张量（对应的n_trg_vocab）还原为训练时的多个特征

import torch.nn as nn
from torch.nn import functional as F

class InputLayer(nn.Module):
    def __init__(self):
        super(InputLayer, self).__init__()
        self.linear1=nn.Linear(4,64)
        self.linear2=nn.Linear(64,1)

    def forward(self,x):
        x = F.relu(self.linear1(x))
        return abs(self.linear2(x))

class OutputLayer(nn.Module):
    def __init__(self):
        super(OutputLayer, self).__init__()
        self.linear1 = nn.Linear(500, 32)
        self.linear2 = nn.Linear(32, 4)

    def forward(self,x):
        x = F.relu(self.linear1(x))
        return self.linear2(x)

训练代码和测试预测代码

# -*- coding: utf-8 -*-
import argparse
import os
import random
import time

import torch
import tqdm
from tensorboardX import SummaryWriter
from torch import optim
from torch.nn import functional as F
from torch.utils.data import DataLoader

from data_read import data_train_raw, WeatherTrajData
from transformer.Models import Transformer
from transformer.Optim import ScheduledOptim

log_writer = SummaryWriter()

os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
print('torch.cuda.is_available=' + str(torch.cuda.is_available()))
torch.set_default_tensor_type(torch.FloatTensor)


# torch.nn.Transformer

def PredictTrajectory(tra_pred, tra_true):

    idx = random.randrange(0, tra_true.shape[0])
    pred = tra_pred[idx, :, :].cpu().detach().numpy()
    true = tra_true[idx, :, :].cpu().detach().numpy()

    print('------------------------------')
    print(f'real value，\n ' +
          f'rho {true[:, 0][-1]}, ' +
          f'sh {true[:, 1][-1]}, ' +
          f'T {true[:, 2][-1]}， ' +
          f'Tdew {true[:, 3][-1]}， ' +
          f'Tlog {true[:, 4][-1]}， ' +
          f'Tpot {true[:, 5][-1]}， ' +
          f'VPact {true[:, 6][-1]}，' +
          f'VPmax {true[:, 7][-1]}')
    # 按照data_read里顺序，下一时刻的位置预测为
    print(f'predict value，\n'
          f'rho {pred[:, 0][-1]}, ' +
          f'sh {pred[:, 1][-1]}, ' +
          f'T {pred[:, 2][-1]}， ' +
          f'Tdew {pred[:, 3][-1]}， ' +
          f'Tlog {pred[:, 4][-1]}， ' +
          f'Tpot {pred[:, 5][-1]}， ' +
          f'VPact {pred[:, 6][-1]}，' +
          f'VPmax {pred[:, 7][-1]}')



def cal_performance(tra_pred, tra_true):
    return F.mse_loss(tra_pred, tra_true)


def train(model, dataloader, optimizer, device, opt):
    for id, epoch_i in enumerate(tqdm.tqdm(range(opt.epoch))):
        model.train()
        total_loss = 0
        for idx, data in enumerate(dataloader):
            optimizer.zero_grad()
            tra_pred = model(input_data=data.to(device).to(torch.float32), device=device)
            # backward and update parameters
            loss = cal_performance(tra_pred, data[:, 1:, :].to(device).to(torch.float32))
            loss.backward()
            optimizer.step_and_update_lr()
            total_loss += loss.item()

        log_writer.add_scalar("loss", total_loss, epoch_i)
        log_writer.add_scalar("lr", optimizer.get_lr(), epoch_i)
        if epoch_i % 100 == 0:
            print("epoch = %d, epoch_loss= %lf ,total_loss = %lf" % (epoch_i, loss.item(), total_loss))

    torch.save(model, 'model.pt')
    # another method to save model
    checkpoint = {
        "net": model.state_dict(),
        "optimizer": optimizer.get_state_dict(),
        "epoch": epoch_i
    }

    if not os.path.isdir("./checkpoint"):
        os.mkdir("./checkpoint")
    torch.save(checkpoint, "./checkpoint/ckpt.pth")

    print("Train Finish")


def test(model, dataloader, device):
    total_loss = 0
    for idx, data in enumerate(dataloader):
        tra_pred = model(input_data=data.to(device).to(torch.float32), device=device)
        loss = cal_performance(tra_pred, data[:, 1:, :].to(device).to(torch.float32))
    total_loss += loss.item()
    PredictTrajectory(tra_pred, data[:, 0:21, :])
    print("Test Finish, total_loss = {}".format(total_loss))


if __name__ == '__main__':
    parser = argparse.ArgumentParser()

    parser.add_argument('-epoch', type=int, default=10000)
    parser.add_argument('-b', '--batch_size', type=int, default=1300)
    parser.add_argument('-d_model', type=int, default=512)
    parser.add_argument('-d_inner_hid', type=int, default=2048)
    parser.add_argument('-d_k', type=int, default=64)
    parser.add_argument('-d_v', type=int, default=64)
    parser.add_argument('-warmup', '--n_warmup_steps', type=int, default=4000)
    parser.add_argument('-lr_mul', type=float, default=2.0)
    parser.add_argument('-lr', type=float, default=0.001)
    parser.add_argument('-n_head', type=int, default=2)
    parser.add_argument('-n_layers', type=int, default=1)
    parser.add_argument('-dropout', type=float, default=0.1)
    parser.add_argument('-do_train', type=bool, default=False)
    parser.add_argument('-do_retrain', type=bool, default=False)
    parser.add_argument('-do_eval', type=bool, default=True)

    opt = parser.parse_args()
    opt.d_word_vec = opt.d_model
    if torch.cuda.is_available():
        # 使用 CUDA 设备
        device = torch.device("cuda")
    else:
        # 使用 CPU 设备
        device = torch.device("cpu")

    print(f"Using device: {device}")

    model_train = Transformer(
        500,
        500,
        d_k=opt.d_k,
        d_v=opt.d_v,
        d_model=opt.d_model,
        d_word_vec=opt.d_word_vec,
        d_inner=opt.d_inner_hid,
        n_layers=opt.n_layers,
        n_head=opt.n_head,
        dropout=opt.dropout,
    ).to(device)

    if opt.do_train == True:
        # data=torch.from_numpy(np.array(data_train)).to(device).to(torch.float32)
        data_train = WeatherTrajData(data_train_raw)
        train_loader = DataLoader(dataset=data_train, batch_size=opt.batch_size, shuffle=False)
        parameters = model_train.parameters()
        optimizer = ScheduledOptim(
            optim.Adam(parameters, betas=(0.9, 0.98), eps=1e-09),
            opt.lr, opt.d_model, opt.n_warmup_steps)

        if opt.do_retrain == True:
            checkpoint = torch.load("./checkpoint/ckpt.pth")
            model_train.load_state_dict(checkpoint['net'])
            optimizer.load_state_dict(checkpoint['optimizer'])

        start_time = time.time()
        train(
            model=model_train,
            dataloader=train_loader,
            optimizer=optimizer,
            device=device,
            opt=opt
        )
        end_time = time.time()
        print("train time = {} seconds".format(end_time - start_time))

    if opt.do_eval == True:
        data_test = WeatherTrajData(data_train_raw)
        test_loader = DataLoader(dataset=data_test, batch_size=opt.batch_size, shuffle=False)
        model = torch.load('model.pt').to(device)
        #model = torch.load('model.pt', map_location=torch.device('cpu'))

        test(
            model=model,
            dataloader=test_loader,
            device=device
        )

开始训练

模型推理

把生成的模型文件下载下来，尝试下预测能力

修改输入参数

parser.add_argument('-do_eval', type=bool, default=True)

比较下真实数据值和预测数据值，还是挺接近的

代码提交在github上

https://github.com/chenrui2200/weather_timeseries_predict