关于人工智能:人工智能创新挑战赛助力精准气象和海洋预测Baseline3TCNNRNN模型SAConvLSTM模型

“AI Earth”人工智能翻新挑战赛：助力精准气象和陆地预测 Baseline[3]：TCNN+RNN 模型、SA-ConvLSTM 模型

1. 气象陆地预测 - 模型建设之 TCNN+RNN

本次工作咱们将学习来自 TOP 选手“swg-lhl”的冠军建模计划，该计划中采纳的模型是 TCNN+RNN。

在 Task3 中咱们学习了 CNN+LSTM 模型，然而 LSTM 层的参数量较大，这就带来以下问题：一是参数量大的模型在数据量小的状况下容易过拟合；二是为了尽量避免过拟合，在无限的数据集下咱们无奈构建更深的模型，难以开掘到更丰盛的信息。相较于 LSTM，CNN 的参数量只与过滤器的大小无关，在各类工作中往往都有不错的体现，因而咱们能够思考同样用卷积操作来开掘工夫信息。然而如果用三维卷积来同时开掘工夫和空间信息，假如应用的过滤器大小为(T_f, H_f, W_f)，那么一层的参数量就是 T_f×H_f×W_f，这样的参数量依然是比拟大的。为了进一步升高每一层的参数，减少模型深度，咱们本次学习的这个 TOP 计划对工夫和空间别离进行卷积操作，即采纳 TCN 单元开掘工夫信息，而后输出 CNN 单元中开掘空间信息，将 TCN 单元 +CNN 单元的串行构造称为 TCNN 层，通过重叠多层的 TCNN 层就能够交替地提取工夫和空间信息。同时，思考到不同时间尺度下的时空信息对预测后果的影响可能是不同的，该计划采纳了三个 RNN 层来抽取三种时间尺度下的特色，将三者拼接起来通过全连贯层预测 Nino3.4 指数。

1.1 数据处理

该 TOP 计划的数据处理次要包含三局部：

1. 数据扁平化。
2. 空值填充。
3. 结构数据集

我的项目链接以及码源见文末

在该计划中除了没有结构新的特色外，其余数据处理办法都与 Task3 基本相同，因而不多做赘述。

import netCDF4 as nc
import random
import os
from tqdm import tqdm
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
%matplotlib inline

import torch
from torch import nn, optim
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader

from sklearn.metrics import mean_squared_error

# 固定随机种子
SEED = 22

def seed_everything(seed=42):
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.backends.cudnn.deterministic = True
    
seed_everything(SEED)

# 查看 GPU 是否可用
train_on_gpu = torch.cuda.is_available()

if not train_on_gpu:
    print('CUDA is not available.  Training on CPU ...')
else:
    print('CUDA is available!  Training on GPU ...')

CUDA is available!  Training on GPU ...

# 读取数据

# 存放数据的门路
path = '/kaggle/input/ninoprediction/'
soda_train = nc.Dataset(path + 'SODA_train.nc')
soda_label = nc.Dataset(path + 'SODA_label.nc')
cmip_train = nc.Dataset(path + 'CMIP_train.nc')
cmip_label = nc.Dataset(path + 'CMIP_label.nc')

1.2 数据扁平化

采纳滑窗结构数据集。

def make_flatted(train_ds, label_ds, info, start_idx=0):
    keys = ['sst', 't300', 'ua', 'va']
    label_key = 'nino'
    # 年数
    years = info[1]
    # 模式数
    models = info[2]
    
    train_list = []
    label_list = []
    
    # 将同种模式下的数据拼接起来
    for model_i in range(models):
        blocks = []
        
        # 对每个特色，取每条数据的前 12 个月进行拼接
        for key in keys:
            block = train_ds[key][start_idx + model_i * years: start_idx + (model_i + 1) * years, :12].reshape(-1, 24, 72, 1).data
            blocks.append(block)
        
        # 将所有特色在最初一个维度上拼接起来
        train_flatted = np.concatenate(blocks, axis=-1)
        
        # 取 12-23 月的标签进行拼接，留神加上最初一年的最初 12 个月的标签（与最初一年 12-23 月的标签独特形成最初一年前 12 个月的预测指标）label_flatted = np.concatenate([label_ds[label_key][start_idx + model_i * years: start_idx + (model_i + 1) * years, 12: 24].reshape(-1).data,
            label_ds[label_key][start_idx + (model_i + 1) * years - 1, 24: 36].reshape(-1).data
        ], axis=0)
        
        train_list.append(train_flatted)
        label_list.append(label_flatted)
        
    return train_list, label_list

soda_info = ('soda', 100, 1)
cmip6_info = ('cmip6', 151, 15)
cmip5_info = ('cmip5', 140, 17)

soda_trains, soda_labels = make_flatted(soda_train, soda_label, soda_info)
cmip6_trains, cmip6_labels = make_flatted(cmip_train, cmip_label, cmip6_info)
cmip5_trains, cmip5_labels = make_flatted(cmip_train, cmip_label, cmip5_info, cmip6_info[1]*cmip6_info[2])

# 失去扁平化后的数据维度为（模式数×序列长度×纬度×经度×特色数），其中序列长度 = 年数×12
np.shape(soda_trains), np.shape(cmip6_trains), np.shape(cmip5_trains)

((1, 1200, 24, 72, 4), (15, 1812, 24, 72, 4), (17, 1680, 24, 72, 4))

1.3 空值填充

将空值填充为 0。

# 填充 SODA 数据中的空值
soda_trains = np.array(soda_trains)
soda_trains_nan = np.isnan(soda_trains)
soda_trains[soda_trains_nan] = 0
print('Number of null in soda_trains after fillna:', np.sum(np.isnan(soda_trains)))

Number of null in soda_trains after fillna: 0

# 填充 CMIP6 数据中的空值
cmip6_trains = np.array(cmip6_trains)
cmip6_trains_nan = np.isnan(cmip6_trains)
cmip6_trains[cmip6_trains_nan] = 0
print('Number of null in cmip6_trains after fillna:', np.sum(np.isnan(cmip6_trains)))

Number of null in cmip6_trains after fillna: 0

# 填充 CMIP5 数据中的空值
cmip5_trains = np.array(cmip5_trains)
cmip5_trains_nan = np.isnan(cmip5_trains)
cmip5_trains[cmip5_trains_nan] = 0
print('Number of null in cmip6_trains after fillna:', np.sum(np.isnan(cmip5_trains)))

Number of null in cmip6_trains after fillna: 0

• 结构数据集
  结构训练和验证集。

# 结构训练集

X_train = []
y_train = []
# 从 CMIP5 的 17 种模式中各抽取 100 条数据
for model_i in range(17):
    samples = np.random.choice(cmip5_trains.shape[1]-12, size=100)
    for ind in samples:
        X_train.append(cmip5_trains[model_i, ind: ind+12])
        y_train.append(cmip5_labels[model_i][ind: ind+24])
# 从 CMIP6 的 15 种模式种各抽取 100 条数据
for model_i in range(15):
    samples = np.random.choice(cmip6_trains.shape[1]-12, size=100)
    for ind in samples:
        X_train.append(cmip6_trains[model_i, ind: ind+12])
        y_train.append(cmip6_labels[model_i][ind: ind+24])
X_train = np.array(X_train)
y_train = np.array(y_train)

# 结构验证集

X_valid = []
y_valid = []
samples = np.random.choice(soda_trains.shape[1]-12, size=100)
for ind in samples:
    X_valid.append(soda_trains[0, ind: ind+12])
    y_valid.append(soda_labels[0][ind: ind+24])
X_valid = np.array(X_valid)
y_valid = np.array(y_valid)

# 查看数据集维度
X_train.shape, y_train.shape, X_valid.shape, y_valid.shape

((3200, 12, 24, 72, 4), (3200, 24), (100, 12, 24, 72, 4), (100, 24))

# 保留数据集
np.save('X_train_sample.npy', X_train)
np.save('y_train_sample.npy', y_train)
np.save('X_valid_sample.npy', X_valid)
np.save('y_valid_sample.npy', y_valid)

1.4 模型构建

# 读取数据集
X_train = np.load('../input/ai-earth-task04-samples/X_train_sample.npy')
y_train = np.load('../input/ai-earth-task04-samples/y_train_sample.npy')
X_valid = np.load('../input/ai-earth-task04-samples/X_valid_sample.npy')
y_valid = np.load('../input/ai-earth-task04-samples/y_valid_sample.npy')

X_train.shape, y_train.shape, X_valid.shape, y_valid.shape

((3200, 12, 24, 72, 4), (3200, 24), (100, 12, 24, 72, 4), (100, 24))

# 结构数据管道
class AIEarthDataset(Dataset):
    def __init__(self, data, label):
        self.data = torch.tensor(data, dtype=torch.float32)
        self.label = torch.tensor(label, dtype=torch.float32)

    def __len__(self):
        return len(self.label)
    
    def __getitem__(self, idx):
        return self.data[idx], self.label[idx]

batch_size = 32

trainset = AIEarthDataset(X_train, y_train)
trainloader = DataLoader(trainset, batch_size=batch_size, shuffle=True)

validset = AIEarthDataset(X_valid, y_valid)
validloader = DataLoader(validset, batch_size=batch_size, shuffle=True)

1.4.1 结构评估函数

def rmse(y_true, y_preds):
    return np.sqrt(mean_squared_error(y_pred = y_preds, y_true = y_true))

# 评估函数
def score(y_true, y_preds):
    # 相关性技巧评分
    accskill_score = 0
    # RMSE
    rmse_scores = 0
    a = [1.5] * 4 + [2] * 7 + [3] * 7 + [4] * 6
    y_true_mean = np.mean(y_true, axis=0)
    y_pred_mean = np.mean(y_preds, axis=0)
    for i in range(24):
        fenzi = np.sum((y_true[:, i] - y_true_mean[i]) * (y_preds[:, i] - y_pred_mean[i]))
        fenmu = np.sqrt(np.sum((y_true[:, i] - y_true_mean[i])**2) * np.sum((y_preds[:, i] - y_pred_mean[i])**2))
        cor_i = fenzi / fenmu
        accskill_score += a[i] * np.log(i+1) * cor_i
        rmse_score = rmse(y_true[:, i], y_preds[:, i])
        rmse_scores += rmse_score
    return 2/3.0 * accskill_score - rmse_scores

1.4.2 模型结构

该 TOP 计划采纳 TCN 单元 +CNN 单元串行组成 TCNN 层，通过重叠多层的 TCNN 层来交替地提取工夫和空间信息，并将提取到的时空信息用 RNN 来抽取出三种不同时间尺度的特色表白。

• TCN 单元

TCN 模型全称工夫卷积网络（Temporal Convolutional Network），与 RNN 一样是时序模型。TCN 以 CNN 为根底，为了适应序列问题，它从以下三方面做出了改良：

1. 因果卷积

TCN 解决输出与输入等长的序列问题，它的每一个暗藏层节点数与输出步长是雷同的，并且暗藏层 t 时刻节点的值只依赖于前一层 t 时刻及之前节点的值。也就是说 TCN 通过追溯前因（t 时刻及之前的值）来取得以后后果，称为因果卷积。

2. 扩张卷积

传统 CNN 的感触野受限于卷积核的大小，须要通过减少池化层来取得更大的感触野，然而池化的操作会带来信息的损失。为了解决这个问题，TCN 采纳扩张卷积来增大感触野，获取更长时间的信息。扩张卷积对输出进行距离采样，采样距离由扩张因子 d 管制，公式定义如下：

$$
F(s) = (X * df)(s) = \sum_{i=0}^{k-1} f(i) \times X_{s-di}
$$

其中 X 为以后层的输出，k 为以后层的卷积核大小，s 为以后节点的时刻。也就是说，对于扩张因子为 d、卷积核为 k 的暗藏层，对前一层的输出每 d 个点采样一次，共采样 k 个点作为以后时刻 s 的输出。这样 TCN 的感触野就由卷积核的大小 k 和扩张因子 d 独特决定，能够获取更长时间的依赖信息。

3. 残差连贯

网络的层数越多，所能提取到的特色就越丰盛，但这也会带来梯度隐没或爆炸的问题，目前解决这个问题的一个无效办法就是残差连贯。TCN 的残差模块蕴含两层卷积操作，并且采纳了 WeightNorm 和 Dropout 进行正则化，如下图所示。

总的来说，TCN 是卷积操作在序列问题上的改良，具备 CNN 参数量少的长处，能够搭建更深层的网络，相比于 RNN 不容易存在梯度隐没和爆炸的问题，同时 TCN 具备灵便的感触野，可能适应不同的工作，在许多数据集上的比拟表明 TCN 比 RNN、LSTM、GRU 等序列模型有更好的体现。

想要更深刻地理解 TCN 能够参考以下链接：

• 论文原文：https://arxiv.org/pdf/1803.01271.pdf
• GitHub：https://github.com/locuslab/tcn

该计划中所构建的 TCN 单元并不是规范的 TCN 层，它的构造如下图所示，能够看到，这里的 TCN 单元只是用了一个卷积层，并且在卷积层前后都采纳了 BatchNormalization 来进步模型的泛化能力。须要留神的是，这里的卷积操作是对工夫维度进行操作，因而须要对输出的形态进行转换，并且为了便于匹配之后的网络层，须要将输入的形态转换回输出时的（N,T,C,H,W）的模式。

# 构建 TCN 单元
class TCNBlock(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding):
        super().__init__()
        self.bn1 = nn.BatchNorm1d(in_channels)
        self.conv = nn.Conv1d(in_channels, out_channels, kernel_size, stride, padding)
        self.bn2 = nn.BatchNorm1d(out_channels)
        
        if in_channels == out_channels and stride == 1:
            self.res = lambda x: x
        else:
            self.res = nn.Conv1d(in_channels, out_channels, kernel_size=1, stride=stride)
        
    def forward(self, x):
        # 转换输出形态
        N, T, C, H, W = x.shape
        x = x.permute(0, 3, 4, 2, 1).contiguous()
        x = x.view(N*H*W, C, T)
        
        # 残差
        res = self.res(x) 
        res = self.bn2(res)

        x = F.relu(self.bn1(x))
        x = self.conv(x)
        x = self.bn2(x)
        
        x = x + res
        
        # 将输入转换回 (N,T,C,H,W) 的模式
        _, C_new, T_new = x.shape
        x = x.view(N, H, W, C_new, T_new)
        x = x.permute(0, 4, 3, 1, 2).contiguous()
        
        return x

• CNN 单元

CNN 单元构造与 TCN 单元类似，都只有一个卷积层，并且应用 BatchNormalization 来进步模型泛化能力。同时，相似 TCN 单元，CNN 单元中也退出了残差连贯。构造如下图所示：

# 构建 CNN 单元
class CNNBlock(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride, padding):
        super().__init__()
        self.bn1 = nn.BatchNorm2d(in_channels)
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding)
        self.bn2 = nn.BatchNorm2d(out_channels)
        
        if (in_channels == out_channels) and (stride == 1):
            self.res = lambda x: x
        else:
            self.res = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=stride)
        
    def forward(self, x):
        # 转换输出形态
        N, T, C, H, W = x.shape
        x = x.view(N*T, C, H, W)
        
        # 残差
        res = self.res(x)
        res = self.bn2(res)

        x = F.relu(self.bn1(x))
        x = self.conv(x)
        x = self.bn2(x)
        
        x = x + res
        
        # 将输入转换回 (N,T,C,H,W) 的模式
        _, C_new, H_new, W_new = x.shape
        x = x.view(N, T, C_new, H_new, W_new)
        
        return x

• TCNN 层

将 TCN 单元和 CNN 单元串行连贯，就形成了一个 TCNN 层。

class TCNNBlock(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride_tcn, stride_cnn, padding):
        super().__init__()
        self.tcn = TCNBlock(in_channels, out_channels, kernel_size, stride_tcn, padding)
        self.cnn = CNNBlock(out_channels, out_channels, kernel_size, stride_cnn, padding)
        
    def forward(self, x):
        x = self.tcn(x)
        x = self.cnn(x)
        return x

• TCNN+RNN 模型

整体的模型构造如下图所示：

1. TCNN 局部

TCNN 局部的模型构造相似传统 CNN 的构造，十分规整，通过逐步减少通道数来提取更丰盛的特色表白。须要留神的是输出数据的格局是(N,T,H,W,C)，为了匹配卷积层的输出格局，须要将数据格式转换为(N,T,C,H,W)。

2. GAP 层

GAP 全称为全局均匀池化（Global Average Pooling）层，它的作用是把每个通道上的特色图取全局均匀，假如通过 TCNN 局部失去的输入格局为(N,T,C,H,W)，那么 GAP 层就会把每个通道上形态为 H×W 的特色图上的所有值求均匀，最终失去的输入格局就变成(N,T,C)。GAP 层最早呈现在论文《Network in Network》（论文原文：https://arxiv.org/pdf/1312.4400.pdf）中用于代替传统 CNN 中的全连贯层，之后的许多试验证实 GAP 层的确能够进步 CNN 的成果。

那么 GAP 层为什么能够代替全连贯层呢？在传统 CNN 中，通过多层卷积和池化的操作后，会由 Flatten 层将特色图拉伸成一列，而后通过全连贯层，那么对于形态为 (C,H,W) 的一条数据，经 Flatten 层拉伸后的长度为 C×H×W，此时假如全连贯层节点数为 U，全连贯层的参数量就是 C×H×W×U，这么大的参数量很容易使得模型过拟合。相比之下，GAP 层不引入新的参数，因而能够无效缩小过拟合问题，并且模型参数少也能放慢训练速度。另一方面，全连贯层是一个黑箱子，咱们很难解释多分类的信息是怎么传回卷积层的，而 GAP 层就很容易了解，每个通道的值就代表了通过多层卷积操作后所提取进去的特色。更具体的了解能够参考 https://www.zhihu.com/question/373188099

在 Pytorch 中没有内置的 GAP 层，因而能够用 adaptive_avg_pool2d 来代替，这个函数能够将特色图压缩成给定的输入形态，将 output_size 参数设置为(1,1)，就等同于 GAP 操作，函数的具体应用办法能够参考 https://pytorch.org/docs/stable/generated/torch.nn.functional…

3. RNN 局部

至此为止咱们所应用的都是长度为 12 的工夫序列，每个工夫步代表一个月的信息。不同尺度的工夫序列所携带的信息是不尽相同的，比方用长度为 6 的工夫序列来表白一年的 SST 值，那么每个工夫步所代表的就是两个月的 SST 信息，这种时间尺度下的 SST 序列与长度为 12 的 SST 序列所反映的一年中 SST 变化趋势等信息就不完全相同。所以，为了尽可能全面地开掘更多信息，该 TOP 计划中用 MaxPool 层来取得三种不同时间尺度的序列，同时，用 RNN 层来抽取序列的特色表白。RNN 非常适合用于线性序列的主动特征提取，例如对于形态为 (T,C1) 的一条输出数据，R 通过节点数为 C2 的 RNN 层就能抽取出长度为 C2 的向量，因为 RNN 由返回后进行信息线性传递的网络结构，抽取出的向量可能很好地表白序列中的依赖关系。

此时三种不同时间尺度的序列都抽取出了一个向量来示意，将向量拼接起来再通过一个全连贯层就失去了 24 个月的预测序列。

# 结构模型
class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Conv2d(4, 64, kernel_size=7, stride=2, padding=3)
        self.tcnn1 = TCNNBlock(64, 64, 3, 1, 1, 1)
        self.tcnn2 = TCNNBlock(64, 128, 3, 1, 2, 1)
        self.tcnn3 = TCNNBlock(128, 128, 3, 1, 1, 1)
        self.tcnn4 = TCNNBlock(128, 256, 3, 1, 2, 1)
        self.tcnn5 = TCNNBlock(256, 256, 3, 1, 1, 1)
        self.rnn = nn.RNN(256, 256, batch_first=True)
        self.maxpool = nn.MaxPool1d(2)
        self.fc = nn.Linear(256*3, 24)
        
    def forward(self, x):
        # 转换输出形态
        N, T, H, W, C = x.shape
        x = x.permute(0, 1, 4, 2, 3).contiguous()
        x = x.view(N*T, C, H, W)
        
        # 通过一个卷积层
        x = self.conv(x)
        _, C_new, H_new, W_new = x.shape
        x = x.view(N, T, C_new, H_new, W_new)
        
        # TCNN 局部
        for i in range(3):
            x = self.tcnn1(x)
        x = self.tcnn2(x)
        for i in range(2):
            x = self.tcnn3(x)
        x = self.tcnn4(x)
        for i in range(2):
            x = self.tcnn5(x)
            
        # 全局均匀池化
        x = F.adaptive_avg_pool2d(x, (1, 1)).squeeze()
        
        # RNN 局部，别离失去长度为 T、T/2、T/ 4 三种时间尺度的特色表白，留神转换 RNN 层输入的格局
        hidden_state = []
        for i in range(3):
            x, h = self.rnn(x)
            h = h.squeeze()
            hidden_state.append(h)
            x = self.maxpool(x.transpose(1, 2)).transpose(1, 2)
        
        x = torch.cat(hidden_state, dim=1)
        x = self.fc(x)
        
        return x

model = Model()
print(model)

1.4.3 模型训练

# 采纳 RMSE 作为损失函数
def RMSELoss(y_pred,y_true):
    loss = torch.sqrt(torch.mean((y_pred-y_true)**2, dim=0)).sum()
    return loss

model_weights = './task04_model_weights.pth'
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = Model().to(device)
criterion = RMSELoss
optimizer = optim.Adam(model.parameters(), lr=1e-3)
epochs = 10
train_losses, valid_losses = [], []
scores = []
best_score = float('-inf')
preds = np.zeros((len(y_valid),24))

for epoch in range(epochs):
    print('Epoch: {}/{}'.format(epoch+1, epochs))
    
    # 模型训练
    model.train()
    losses = 0
    for data, labels in tqdm(trainloader):
        data = data.to(device)
        labels = labels.to(device)
        optimizer.zero_grad()
        pred = model(data)
        loss = criterion(pred, labels)
        losses += loss.cpu().detach().numpy()
        loss.backward()
        optimizer.step()
    train_loss = losses / len(trainloader)
    train_losses.append(train_loss)
    print('Training Loss: {:.3f}'.format(train_loss))
    
    # 模型验证
    model.eval()
    losses = 0
    with torch.no_grad():
        for i, data in tqdm(enumerate(validloader)):
            data, labels = data
            data = data.to(device)
            labels = labels.to(device)
            pred = model(data)
            loss = criterion(pred, labels)
            losses += loss.cpu().detach().numpy()
            preds[i*batch_size:(i+1)*batch_size] = pred.detach().cpu().numpy()
    valid_loss = losses / len(validloader)
    valid_losses.append(valid_loss)
    print('Validation Loss: {:.3f}'.format(valid_loss))
    s = score(y_valid, preds)
    scores.append(s)
    print('Score: {:.3f}'.format(s))
    
    # 保留最佳模型权重
    if s > best_score:
        best_score = s
        checkpoint = {'best_score': s,
                      'state_dict': model.state_dict()}
        torch.save(checkpoint, model_weights)

# 绘制训练 / 验证曲线
def training_vis(train_losses, valid_losses):
    # 绘制损失函数曲线
    fig = plt.figure(figsize=(8,4))
    # subplot loss
    ax1 = fig.add_subplot(121)
    ax1.plot(train_losses, label='train_loss')
    ax1.plot(valid_losses,label='val_loss')
    ax1.set_xlabel('Epochs')
    ax1.set_ylabel('Loss')
    ax1.set_title('Loss on Training and Validation Data')
    ax1.legend()
    plt.tight_layout()

training_vis(train_losses, valid_losses)

1.4.4 模型评估

在测试集上评估模型成果。

# 加载最佳模型权重
checkpoint = torch.load('../input/ai-earth-model-weights/task04_model_weights.pth')
model = Model()
model.load_state_dict(checkpoint['state_dict'])

# 测试集门路
test_path = '../input/ai-earth-tests/'
# 测试集标签门路
test_label_path = '../input/ai-earth-tests-labels/'

import os

# 读取测试数据和测试数据的标签
files = os.listdir(test_path)
X_test = []
y_test = []
for file in files:
    X_test.append(np.load(test_path + file))
    y_test.append(np.load(test_label_path + file))

X_test = np.array(X_test)
y_test = np.array(y_test)
X_test.shape, y_test.shape

((103, 12, 24, 72, 4), (103, 24))

testset = AIEarthDataset(X_test, y_test)
testloader = DataLoader(testset, batch_size=batch_size, shuffle=False)

# 在测试集上评估模型成果
model.eval()
model.to(device)
preds = np.zeros((len(y_test),24))
for i, data in tqdm(enumerate(testloader)):
    data, labels = data
    data = data.to(device)
    labels = labels.to(device)
    pred = model(data)
    preds[i*batch_size:(i+1)*batch_size] = pred.detach().cpu().numpy()
s = score(y_test, preds)
print('Score: {:.3f}'.format(s))

4it [00:00, 12.75it/s]

Score: 20.274

1.5 小结

• 该计划充分考虑到数据量小、特色少的数据状况，对工夫和空间别离进行卷积操作，交替地提取工夫和空间信息，用 GAP 层对提取的信息进行降维，尽可能减少每一层的参数量、减少模型层数以提取更丰盛的特色。
• 该计划思考到不同时间尺度序列所携带的信息不同，用池化层变换时间尺度，并用 RNN 进行信息提取，综合三种不同时间尺度的序列信息失去最终的预测序列。
• 该计划同样抉择了本人设计模型，在结构模型时充分考虑了数据集状况和问题背景，并能灵活运用各种网络层来解决特定问题，这种模型结构思路要求对不同网络层的作用有较为粗浅地了解，计划中各种网络层的用法值得大家学习和借鉴。

参考文献

1. Top1 思路分享：https://tianchi.aliyun.com/forum/postDetail?spm=5176.12586969…

2. 气象陆地预测 - 模型建设之 SA-ConvLSTM

本次工作咱们将学习来自 TOP 选手“吴先生的队伍”的建模计划，该计划中采纳的模型是 SA-ConvLSTM。

前两个 TOP 计划中抉择将赛题看作一个多输入的工作，通过构建神经网络间接输入 24 个 nino3.4 预测值，这种思路的问题在于，序列问题往往是时序依赖的，当咱们采纳多输入的办法时其实把这 24 个 nino3.4 预测值看作是齐全独立的，然而实际上它们之间是存在序列依赖的，即每个预测值往往受上一个工夫步的预测值的影响。因而，在这次的 TOP 计划中，采纳 Seq2Seq 构造来思考输入预测值的序列依赖性。

Seq2Seq 构造包含 Encoder（编码器）和 Decoder（解码器）两局部，Encoder 局部将输出序列编码成一个向量，Decoder 局部对向量进行解码，输入一个预测序列。要将 Seq2Seq 构造利用于不同的序列问题，关键在于每一个工夫步所应用的 Cell。咱们之前说到，开掘空间信息通常会采纳 CNN，开掘工夫信息通常会采纳 RNN 或 LSTM，将二者联合在一起就失去了时空序列畛域的经典模型——ConvLSTM，咱们本次要学习的 SA-ConvLSTM 模型是对 ConvLSTM 模型的改良，在其根底上引入了自注意力机制来进步模型对于长期空间依赖关系的开掘能力。

另外与前两个 TOP 计划所不同的一点是，该 TOP 计划没有间接预测 Nino3.4 指数，而是通过预测 sst 来间接求得 Nino3.4 指数序列。

import netCDF4 as nc
import random
import os
from tqdm import tqdm
import pandas as pd
import numpy as np
import math
import matplotlib.pyplot as plt
%matplotlib inline

import torch
from torch import nn, optim
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from torch.optim.lr_scheduler import ReduceLROnPlateau

from sklearn.metrics import mean_squared_error

# 固定随机种子
SEED = 22

def seed_everything(seed=42):
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.backends.cudnn.deterministic = True
    
seed_everything(SEED)

# 查看 CUDA 是否可用
train_on_gpu = torch.cuda.is_available()

if not train_on_gpu:
    print('CUDA is not available.  Training on CPU ...')
else:
    print('CUDA is available!  Training on GPU ...')

CUDA is available!  Training on GPU ...

# 读取数据

# 存放数据的门路
path = '/kaggle/input/ninoprediction/'
soda_train = nc.Dataset(path + 'SODA_train.nc')
soda_label = nc.Dataset(path + 'SODA_label.nc')
cmip_train = nc.Dataset(path + 'CMIP_train.nc')
cmip_label = nc.Dataset(path + 'CMIP_label.nc')

2.1 数据扁平化

采纳滑窗结构数据集。该计划中只应用了 sst 特色，且只应用了 lon 值在 [90, 330] 范畴内的数据，可能是为了节约计算资源。

def make_flatted(train_ds, label_ds, info, start_idx=0):
    # 只应用 sst 特色
    keys = ['sst']
    label_key = 'nino'
    # 年数
    years = info[1]
    # 模式数
    models = info[2]
    
    train_list = []
    label_list = []
    
    # 将同种模式下的数据拼接起来
    for model_i in range(models):
        blocks = []
        
        # 对每个特色，取每条数据的前 12 个月进行拼接，只应用 lon 值在 [90, 330] 范畴内的数据
        for key in keys:
            block = train_ds[key][start_idx + model_i * years: start_idx + (model_i + 1) * years, :12, :, 19: 67].reshape(-1, 24, 48, 1).data
            blocks.append(block)
        
        # 将所有特色在最初一个维度上拼接起来
        train_flatted = np.concatenate(blocks, axis=-1)
        
        # 取 12-23 月的标签进行拼接，留神加上最初一年的最初 12 个月的标签（与最初一年 12-23 月的标签独特形成最初一年前 12 个月的预测指标）label_flatted = np.concatenate([label_ds[label_key][start_idx + model_i * years: start_idx + (model_i + 1) * years, 12: 24].reshape(-1).data,
            label_ds[label_key][start_idx + (model_i + 1) * years - 1, 24: 36].reshape(-1).data
        ], axis=0)
        
        train_list.append(train_flatted)
        label_list.append(label_flatted)
        
    return train_list, label_list

soda_info = ('soda', 100, 1)
cmip6_info = ('cmip6', 151, 15)
cmip5_info = ('cmip5', 140, 17)

soda_trains, soda_labels = make_flatted(soda_train, soda_label, soda_info)
cmip6_trains, cmip6_labels = make_flatted(cmip_train, cmip_label, cmip6_info)
cmip5_trains, cmip5_labels = make_flatted(cmip_train, cmip_label, cmip5_info, cmip6_info[1]*cmip6_info[2])

# 失去扁平化后的数据维度为（模式数×序列长度×纬度×经度×特色数），其中序列长度 = 年数×12
np.shape(soda_trains), np.shape(cmip6_trains), np.shape(cmip5_trains)

((1, 1200, 24, 48, 1), (15, 1812, 24, 48, 1), (17, 1680, 24, 48, 1))

空值填充
将空值填充为 0。

# 填充 SODA 数据中的空值
soda_trains = np.array(soda_trains)
soda_trains_nan = np.isnan(soda_trains)
soda_trains[soda_trains_nan] = 0
print('Number of null in soda_trains after fillna:', np.sum(np.isnan(soda_trains)))

Number of null in soda_trains after fillna: 0

# 填充 CMIP6 数据中的空值
cmip6_trains = np.array(cmip6_trains)
cmip6_trains_nan = np.isnan(cmip6_trains)
cmip6_trains[cmip6_trains_nan] = 0
print('Number of null in cmip6_trains after fillna:', np.sum(np.isnan(cmip6_trains)))

Number of null in cmip6_trains after fillna: 0

# 填充 CMIP5 数据中的空值
cmip5_trains = np.array(cmip5_trains)
cmip5_trains_nan = np.isnan(cmip5_trains)
cmip5_trains[cmip5_trains_nan] = 0
print('Number of null in cmip6_trains after fillna:', np.sum(np.isnan(cmip5_trains)))

Number of null in cmip6_trains after fillna: 0

2.2 结构数据集

结构训练和验证集。留神这里取每条输出数据的序列长度是 38，这是因为输出 sst 序列长度是 12，输入 sst 序列长度是 26，在训练中采纳 teacher forcing 策略（这个策略会在之后的模型结构时具体阐明），因而这里在结构输出数据时蕴含了输入 sst 序列的理论值。

# 结构训练集

X_train = []
y_train = []
# 从 CMIP5 的 17 种模式中各抽取 100 条数据
for model_i in range(17):
    samples = np.random.choice(cmip5_trains.shape[1]-38, size=100)
    for ind in samples:
        X_train.append(cmip5_trains[model_i, ind: ind+38])
        y_train.append(cmip5_labels[model_i][ind: ind+24])
# 从 CMIP6 的 15 种模式种各抽取 100 条数据
for model_i in range(15):
    samples = np.random.choice(cmip6_trains.shape[1]-38, size=100)
    for ind in samples:
        X_train.append(cmip6_trains[model_i, ind: ind+38])
        y_train.append(cmip6_labels[model_i][ind: ind+24])
X_train = np.array(X_train)
y_train = np.array(y_train)

# 结构测试集

X_valid = []
y_valid = []
samples = np.random.choice(soda_trains.shape[1]-38, size=100)
for ind in samples:
    X_valid.append(soda_trains[0, ind: ind+38])
    y_valid.append(soda_labels[0][ind: ind+24])
X_valid = np.array(X_valid)
y_valid = np.array(y_valid)

# 查看数据集维度
X_train.shape, y_train.shape, X_valid.shape, y_valid.shape

((3200, 38, 24, 48, 1), (3200, 24), (100, 38, 24, 48, 1), (100, 24))

# 保留数据集
np.save('X_train_sample.npy', X_train)
np.save('y_train_sample.npy', y_train)
np.save('X_valid_sample.npy', X_valid)
np.save('y_valid_sample.npy', y_valid)

2.3 模型构建

# 读取数据集
X_train = np.load('../input/ai-earth-task05-samples/X_train_sample.npy')
y_train = np.load('../input/ai-earth-task05-samples/y_train_sample.npy')
X_valid = np.load('../input/ai-earth-task05-samples/X_valid_sample.npy')
y_valid = np.load('../input/ai-earth-task05-samples/y_valid_sample.npy')

X_train.shape, y_train.shape, X_valid.shape, y_valid.shape

((3200, 38, 24, 48, 1), (3200, 24), (100, 38, 24, 48, 1), (100, 24))

# 结构数据管道
class AIEarthDataset(Dataset):
    def __init__(self, data, label):
        self.data = torch.tensor(data, dtype=torch.float32)
        self.label = torch.tensor(label, dtype=torch.float32)

    def __len__(self):
        return len(self.label)
    
    def __getitem__(self, idx):
        return self.data[idx], self.label[idx]

batch_size = 2

trainset = AIEarthDataset(X_train, y_train)
trainloader = DataLoader(trainset, batch_size=batch_size, shuffle=True)

validset = AIEarthDataset(X_valid, y_valid)
validloader = DataLoader(validset, batch_size=batch_size, shuffle=True)

2.3.1 结构评估函数

def rmse(y_true, y_preds):
    return np.sqrt(mean_squared_error(y_pred = y_preds, y_true = y_true))

# 评估函数
def score(y_true, y_preds):
    # 相关性技巧评分
    accskill_score = 0
    # RMSE
    rmse_scores = 0
    a = [1.5] * 4 + [2] * 7 + [3] * 7 + [4] * 6
    y_true_mean = np.mean(y_true, axis=0)
    y_pred_mean = np.mean(y_preds, axis=0)
    for i in range(24):
        fenzi = np.sum((y_true[:, i] - y_true_mean[i]) * (y_preds[:, i] - y_pred_mean[i]))
        fenmu = np.sqrt(np.sum((y_true[:, i] - y_true_mean[i])**2) * np.sum((y_preds[:, i] - y_pred_mean[i])**2))
        cor_i = fenzi / fenmu
        accskill_score += a[i] * np.log(i+1) * cor_i
        rmse_score = rmse(y_true[:, i], y_preds[:, i])
        rmse_scores += rmse_score
    return 2/3.0 * accskill_score - rmse_scores

2.3.2 模型结构

不同于前两个 TOP 计划所构建的多输入神经网络，该 TOP 计划采纳的是 Seq2Seq 构造，以本赛题为例，输出的序列长度是 12，输入的序列长度是 26，计划中构建了四个暗藏层，那么一个根底的 Seq2Seq 构造就如下图所示：

要将 Seq2Seq 构造利用于不同的问题，重点在于应用怎么的 Cell（神经元）。在该 TOP 计划中应用的 Cell 是清华大学提出的 SA-ConvLSTM（Self-Attention ConvLSTM），论文原文可参考 https://ojs.aaai.org//index.php/AAAI/article/view/6819

SA-ConvLSTM 是实施健博士提出的时空序列畛域经典模型 ConvLSTM 的改良模型，为了捕获空间信息的时序依赖关系，它在 ConvLSTM 的根底上减少了 SAM 模块，用来记忆空间的聚合特色。ConvLSTM 的论文原文可参考 https://arxiv.org/pdf/1506.04214.pdf

1. ConvLSTM 模型

LSTM 模型是十分经典的时序模型，三个门的构造使得它在开掘长期的工夫依赖工作中有不俗的体现，并且相较于 RNN，LSTM 可能无效地防止梯度隐没问题。对于单个输出样本，在每个工夫步上，LSTM 的每个门理论是对输出向量做了一个全连贯，那么对应到咱们这个赛题上，输出 X 的形态是（N,T,H,W,C），则单个输出样本在每个工夫步上输出 LSTM 的就是形态为（H,W,C）的空间信息。咱们晓得，全连贯网络对于这种空间信息的提取能力并不强，转换成卷积操作后可能在大大减少参数量的同时通过重叠多层网络逐渐提取出更简单的特色，到这里就能够很天然地想到，把 LSTM 中的全连贯操作转换为卷积操作，就可能实用于时空序列问题。ConvLSTM 模型就是这么做的，实际也表明这样的作法是十分无效的。

2. SAM 模块

然而，ConvLSTM 模型存在两个问题：

一是卷积层的感触野受限于卷积核的大小，须要通过重叠多个卷积层来扩充感触野，挖掘全局的特色。举例来说，假如第一个卷积层的卷积核大小是 3×3，那么这一层的每个节点就只能感知这 3×3 的空间范畴内的输出信息，此时再减少一个 3×3 的卷积层，那么每个节点所能感知的就是 3×3 个第一层的节点内的信息，在第一层步长为 1 的状况下，就是 4×4 范畴内的输出信息，于是相比于第一个卷积层，第二层所能感知的输出信息的空间范畴就增大了，而这样做所带来的结果就是参数量减少。对于单纯的 CNN 模型来说减少一层只是减少了一个卷积核大小的参数量，然而对于 ConvLSTM 来说就有些不堪重负，参数量的减少增大了过拟合的危险，与此同时模型的收效却并不高。

二是卷积操作只针对以后工夫步输出的空间信息，而漠视了过来的空间信息，因而难以开掘空间信息在工夫上的依赖关系。

因而，为了同时开掘全局和本地的空间依赖，晋升模型在大空间范畴和长时间的时空序列预测工作中的预测成果，SA-ConvLSTM 模型在 ConvLSTM 模型的根底上引入了 SAM（self-attention memory）模块。

SAM 模块引入了一个新的记忆单元 M，用来记忆蕴含时序依赖关系的空间信息。SAM 模块以以后工夫步通过 ConvLSTM 所取得的暗藏层状态 $H_t$ 和上一个工夫步的记忆 $M_{t-1}$ 作为输出，首先将 $H_t$ 通过自注意力机制失去特色 $Z_h$，自注意力机制可能减少 $H_t$ 中与其余局部更相干的局部的权重，同时 $H_t$ 也作为 Query 与 $M_{t-1}$ 独特通过注意力机制失去特色 $Z_m$，用以加强对 $M_{t-1}$ 中与 $H_t$ 有更强依赖关系的局部的权重，将 $Z_h$ 和 $Z_m$ 拼接起来就失去了二者的聚合特色 $Z$。此时，聚合特色 $Z$ 中既蕴含了以后工夫步的信息，又蕴含了全局的时空记忆信息，接下来借鉴 LSTM 中的门控构造用聚合特色 $Z$ 对暗藏层状态和记忆单元进行更新，就失去了更新后的暗藏层状态 $\hat{H_t}$ 和以后工夫步的记忆 $M_t$。SAM 模块的公式如下：

$$
\begin{aligned}
& i’_t = \sigma (W_{m;zi} \ast Z + W_{m;hi} \ast H_t + b_{m;i}) \\
& g’_t = tanh (W_{m;zg} \ast Z + W_{m;hg} \ast H_t + b_{m;g}) \\
& M_t = (1 – i’_t) \circ M_{t-1} + i’_t \circ g’_t \\
& o’_t = \sigma (W_{m;zo} \ast Z + W_{m;ho} \ast H_t + b_{m;o}) \\
& \hat{H_t} = o’_t \circ M_t
\end{aligned}
$$

对于注意力机制和自注意力机制能够参考以下链接：

• 深度学习中的注意力机制：https://blog.csdn.net/malefactor/article/details/78767781
• 目前支流的 Attention 办法：https://www.zhihu.com/question/68482809

3. SA-ConvLSTM 模型

将以上二者联合起来，就失去了 SA-ConvLSTM 模型：

# Attention 机制
def attn(query, key, value):
    # query、key、value 的形态都是(N, C, H*W)，令 S =H*W
    # 采纳缩放点积模型计算得分，scores(i)=key(i)^T query/ 根号 C
    scores = torch.matmul(query.transpose(1, 2), key / math.sqrt(query.size(1)))  # (N, S, S)
    # 计算注意力得分
    attn = F.softmax(scores, dim=-1)
    output = torch.matmul(attn, value.transpose(1, 2))  # (N, S, C)
    return output.transpose(1, 2)  # (N, C, S)

# SAM 模块
class SAAttnMem(nn.Module):
    def __init__(self, input_dim, d_model, kernel_size):
        super().__init__()
        pad = kernel_size[0] // 2, kernel_size[1] // 2
        self.d_model = d_model
        self.input_dim = input_dim
        # 用 1 * 1 卷积实现全连贯操作 WhHt
        self.conv_h = nn.Conv2d(input_dim, d_model*3, kernel_size=1)
        # 用 1 * 1 卷积实现全连贯操作 WmMt-1
        self.conv_m = nn.Conv2d(input_dim, d_model*2, kernel_size=1)
        # 用 1 * 1 卷积实现全连贯操作 Wz[Zh,Zm]
        self.conv_z = nn.Conv2d(d_model*2, d_model, kernel_size=1)
        # 留神输入维度和输出维度要保持一致，都是 input_dim
        self.conv_output = nn.Conv2d(input_dim+d_model, input_dim*3, kernel_size=kernel_size, padding=pad)
        
    def forward(self, h, m):
        # self.conv_h(h)失去 WhHt，将其在 dim= 1 上划分成大小为 self.d_model 的块，每一块的形态就是(N, d_model, H, W)，所失去的三块就是 Qh、Kh、Vh
        hq, hk, hv = torch.split(self.conv_h(h), self.d_model, dim=1)
        # 同样的办法失去 Km 和 Vm
        mk, mv = torch.split(self.conv_m(m), self.d_model, dim=1)
        N, C, H, W = hq.size()
        # 通过自注意力机制失去 Zh
        Zh = attn(hq.view(N, C, -1), hk.view(N, C, -1), hv.view(N, C, -1))  # (N, C, S), C=d_model
        # 通过注意力机制失去 Zm
        Zm = attn(hq.view(N, C, -1), mk.view(N, C, -1), mv.view(N, C, -1))  # (N, C, S), C=d_model
        # 将 Zh 和 Zm 拼接起来，并进行全连贯操作失去聚合特色 Z
        Z = self.conv_z(torch.cat([Zh.view(N, C, H, W), Zm.view(N, C, H, W)], dim=1))  # (N, C, H, W), C=d_model
        # 计算 i't、g't、o't
        i, g, o = torch.split(self.conv_output(torch.cat([Z, h], dim=1)), self.input_dim, dim=1)  # (N, C, H, W), C=input_dim
        i = torch.sigmoid(i)
        g = torch.tanh(g)
        # 失去更新后的记忆单元 Mt
        m_next = i * g + (1 - i) * m
        # 失去更新后的暗藏状态 Ht
        h_next = torch.sigmoid(o) * m_next
        return h_next, m_next

# SA-ConvLSTM Cell
class SAConvLSTMCell(nn.Module):
    def __init__(self, input_dim, hidden_dim, d_attn, kernel_size):
        super().__init__()
        self.input_dim = input_dim
        self.hidden_dim = hidden_dim
        pad = kernel_size[0] // 2, kernel_size[1] // 2
        # 卷积操作 Wx*Xt+Wh*Ht-1
        self.conv = nn.Conv2d(in_channels=input_dim+hidden_dim, out_channels=4*hidden_dim, kernel_size=kernel_size, padding=pad)
        self.sa = SAAttnMem(input_dim=hidden_dim, d_model=d_attn, kernel_size=kernel_size)
        
    def initialize(self, inputs):
        device = inputs.device
        N, _, H, W = inputs.size()
        # 初始化暗藏层状态 Ht
        self.hidden_state = torch.zeros(N, self.hidden_dim, H, W, device=device)
        # 初始化记忆细胞状态 ct
        self.cell_state = torch.zeros(N, self.hidden_dim, H, W, device=device)
        # 初始化记忆单元状态 Mt
        self.memory_state = torch.zeros(N, self.hidden_dim, H, W, device=device)
    
    def forward(self, inputs, first_step=False):
        # 如果以后是第一个工夫步，初始化 Ht、ct、Mt
        if first_step:
            self.initialize(inputs)
        
        # ConvLSTM 局部
        # 拼接 Xt 和 Ht
        combined = torch.cat([inputs, self.hidden_state], dim=1)  # (N, C, H, W), C=input_dim+hidden_dim
        # 进行卷积操作
        combined_conv = self.conv(combined)       
        # 失去四个门控单元 it、ft、ot、gt
        cc_i, cc_f, cc_o, cc_g = torch.split(combined_conv, self.hidden_dim, dim=1)
        i = torch.sigmoid(cc_i)
        f = torch.sigmoid(cc_f)
        o = torch.sigmoid(cc_o)
        g = torch.tanh(cc_g)
        # 失去以后工夫步的记忆细胞状态 ct=ft·ct-1+it·gt
        self.cell_state = f * self.cell_state + i * g
        # 失去以后工夫步的暗藏层状态 Ht=ot·tanh(ct)
        self.hidden_state = o * torch.tanh(self.cell_state)
        
        # SAM 局部，更新 Ht 和 Mt
        self.hidden_state, self.memory_state = self.sa(self.hidden_state, self.memory_state)
        
        return self.hidden_state

在 Seq2Seq 模型的训练中，有两种训练模式。一是 Free running，也就是传统的训练形式，以上一个工夫步的输入 $\hat{y_{t-1}}$ 作为下一个工夫步的输出，然而这种做法存在的问题是在训练的初期所失去的 $\hat{y_{t-1}}$ 与理论标签 $y_{t-1}$ 相差甚远，以此作为输出会导致后续的输入越来越偏离咱们冀望的预测标签。于是就产生了第二种训练模式——Teacher forcing。

Teacher forcing 就是间接应用理论标签 $y_{t-1}$ 作为下一个工夫步的输出，由老师（ground truth）率领着避免模型越走越偏。然而老师不能总是手把手领着学生走，要逐步撒手让学生自主学习，于是咱们应用 Scheduled Sampling 来管制应用理论标签的概率。咱们用 ratio 来示意 Scheduled Sampling 的比例，在训练初期，ratio=1，模型齐全由老师率领着，随着训练阐述的减少，ratio 以肯定的形式衰减（该计划中应用线性衰减，ratio 每次减小一个衰减率 decay_rate），每个工夫步以 ratio 的概率从伯努利散布中提取二进制随机数 0 或 1，为 1 时输出就是理论标签 $y_{t-1}$，否则输出为 $\hat{y_{t-1}}$。

# 构建 SA-ConvLSTM 模型
class SAConvLSTM(nn.Module):
    def __init__(self, input_dim, hidden_dim, d_attn, kernel_size):
        super().__init__()
        self.input_dim = input_dim
        self.hidden_dim = hidden_dim
        self.num_layers = len(hidden_dim)
        
        layers = []
        for i in range(self.num_layers):
            cur_input_dim = self.input_dim if i == 0 else self.hidden_dim[i-1]
            layers.append(SAConvLSTMCell(input_dim=cur_input_dim, hidden_dim=self.hidden_dim[i], d_attn = d_attn, kernel_size=kernel_size))            
        self.layers = nn.ModuleList(layers)
        
        self.conv_output = nn.Conv2d(self.hidden_dim[-1], 1, kernel_size=1)
        
    def forward(self, input_x, device=torch.device('cuda:0'), input_frames=12, future_frames=26, output_frames=37, teacher_forcing=False, scheduled_sampling_ratio=0, train=True):
        # 将输出样本 X 的形态 (N, T, H, W, C) 转换为(N, T, C, H, W)
        input_x = input_x.permute(0, 1, 4, 2, 3).contiguous()
        
        # 仅在训练时应用 teacher forcing
        if train:
            if teacher_forcing and scheduled_sampling_ratio > 1e-6:
                teacher_forcing_mask = torch.bernoulli(scheduled_sampling_ratio * torch.ones(input_x.size(0), future_frames-1, 1, 1, 1))
            else:
                teacher_forcing = False
        else:
            teacher_forcing = False
            
        total_steps = input_frames + future_frames - 1
        outputs = [None] * total_steps
        
        # 对于每一个工夫步
        for t in range(total_steps):
            # 在前 12 个月，应用每个月的输出样本 Xt
            if t < input_frames:
                input_ = input_x[:, t].to(device)
            # 若不应用 teacher forcing，则以上一个工夫步的预测标签作为以后工夫步的输出
            elif not teacher_forcing:
                input_ = outputs[t-1]
            # 若应用 teacher forcing，则以 ratio 的概率应用上一个工夫步的理论标签作为以后工夫步的输出
            else:
                mask = teacher_forcing_mask[:, t-input_frames].float().to(device)
                input_ = input_x[:, t].to(device) * mask + outputs[t-1] * (1-mask)
            first_step = (t==0)
            input_ = input_.float()
            
            # 将以后工夫步的输出通过暗藏层
            for layer_idx in range(self.num_layers):
                input_ = self.layers[layer_idx](input_, first_step=first_step)
            
            # 记录每个工夫步的输入
            if train or (t >= (input_frames - 1)):
                outputs[t] = self.conv_output(input_)
                
        outputs = [x for x in outputs if x is not None]
        
        # 确认输入序列的长度
        if train:
            assert len(outputs) == output_frames
        else:
            assert len(outputs) == future_frames
        
        # 失去 sst 的预测序列
        outputs = torch.stack(outputs, dim=1)[:, :, 0]  # (N, 37, H, W)
        # 对 sst 的预测序列在 nino3.4 区域取三个月的平均值就失去 nino3.4 指数的预测序列
        nino_pred = outputs[:, -future_frames:, 10:13, 19:30].mean(dim=[2, 3])  # (N, 26)
        nino_pred = nino_pred.unfold(dimension=1, size=3, step=1).mean(dim=2)  # (N, 24)
        
        return nino_pred

# 输出特色数
input_dim = 1
# 暗藏层节点数
hidden_dim = (64, 64, 64, 64)
# 注意力机制节点数
d_attn = 32
# 卷积核大小
kernel_size = (3, 3)

model = SAConvLSTM(input_dim, hidden_dim, d_attn, kernel_size)
print(model)

2.3.3 模型训练

# 采纳 RMSE 作为损失函数
def RMSELoss(y_pred,y_true):
    loss = torch.sqrt(torch.mean((y_pred-y_true)**2, dim=0)).sum()
    return loss

model_weights = './task05_model_weights.pth'
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = SAConvLSTM(input_dim, hidden_dim, d_attn, kernel_size).to(device)
criterion = RMSELoss
optimizer = optim.Adam(model.parameters(), lr=1e-3)
lr_scheduler = ReduceLROnPlateau(optimizer, mode='max', factor=0.3, patience=0, verbose=True, min_lr=0.0001)
epochs = 5
ratio, decay_rate = 1, 8e-5
train_losses, valid_losses = [], []
scores = []
best_score = float('-inf')
preds = np.zeros((len(y_valid),24))

for epoch in range(epochs):
    print('Epoch: {}/{}'.format(epoch+1, epochs))
    
    # 模型训练
    model.train()
    losses = 0
    for data, labels in tqdm(trainloader):
        data = data.to(device)
        labels = labels.to(device)
        optimizer.zero_grad()
        # ratio 线性衰减
        ratio = max(ratio-decay_rate, 0)
        pred = model(data, teacher_forcing=True, scheduled_sampling_ratio=ratio, train=True)
        loss = criterion(pred, labels)
        losses += loss.cpu().detach().numpy()
        loss.backward()
        optimizer.step()
    train_loss = losses / len(trainloader)
    train_losses.append(train_loss)
    print('Training Loss: {:.3f}'.format(train_loss))
    
    # 模型验证
    model.eval()
    losses = 0
    with torch.no_grad():
        for i, data in tqdm(enumerate(validloader)):
            data, labels = data
            data = data.to(device)
            labels = labels.to(device)
            pred = model(data, train=False)
            loss = criterion(pred, labels)
            losses += loss.cpu().detach().numpy()
            preds[i*batch_size:(i+1)*batch_size] = pred.detach().cpu().numpy()
    valid_loss = losses / len(validloader)
    valid_losses.append(valid_loss)
    print('Validation Loss: {:.3f}'.format(valid_loss))
    s = score(y_valid, preds)
    scores.append(s)
    print('Score: {:.3f}'.format(s))
    
    # 保留最佳模型权重
    if s > best_score:
        best_score = s
        checkpoint = {'best_score': s,
                      'state_dict': model.state_dict()}
        torch.save(checkpoint, model_weights)

# 绘制训练 / 验证曲线
def training_vis(train_losses, valid_losses):
    # 绘制损失函数曲线
    fig = plt.figure(figsize=(8,4))
    # subplot loss
    ax1 = fig.add_subplot(121)
    ax1.plot(train_losses, label='train_loss')
    ax1.plot(valid_losses,label='val_loss')
    ax1.set_xlabel('Epochs')
    ax1.set_ylabel('Loss')
    ax1.set_title('Loss on Training and Validation Data')
    ax1.legend()
    plt.tight_layout()

training_vis(train_losses, valid_losses)

2.3.4 模型评估

在测试集上评估模型成果。

# 加载得分最高的模型
checkpoint = torch.load('../input/ai-earth-model-weights/task05_model_weights.pth')
model = SAConvLSTM(input_dim, hidden_dim, d_attn, kernel_size)
model.load_state_dict(checkpoint['state_dict'])

<All keys matched successfully>

# 测试集门路
test_path = '../input/ai-earth-tests/'
# 测试集标签门路
test_label_path = '../input/ai-earth-tests-labels/'

import os

# 读取测试数据和测试数据的标签
files = os.listdir(test_path)
X_test = []
y_test = []
for file in files:
    X_test.append(np.load(test_path + file))
    y_test.append(np.load(test_label_path + file))

X_test = np.array(X_test)[:, :, :, 19: 67, :1]
y_test = np.array(y_test)
X_test.shape, y_test.shape

((103, 12, 24, 48, 1), (103, 24))

testset = AIEarthDataset(X_test, y_test)
testloader = DataLoader(testset, batch_size=batch_size, shuffle=False)

# 在测试集上评估模型成果
model.eval()
model.to(device)
preds = np.zeros((len(y_test),24))
for i, data in tqdm(enumerate(testloader)):
    data, labels = data
    data = data.to(device)
    labels = labels.to(device)
    pred = model(data, train=False)
    preds[i*batch_size:(i+1)*batch_size] = pred.detach().cpu().numpy()
s = score(y_test, preds)
print('Score: {:.3f}'.format(s))

2.4 小结

这一次的 TOP 计划没有本人设计模型，而是应用了目前时空序列预测畛域现有的模型，另一组 TOP 选手“ailab”也应用了现有的模型 PredRNN++，对于时空序列预测畛域的一些比拟经典的模型能够参考 https://www.zhihu.com/column/c_1208033701705162752

我的项目链接以及码源见文末

链接跳转到文末，欢送订阅

更多文章请关注公重号：汀丶人工智能