气象海洋预测-Task5 模型建立之 SA-ConvLSTM

该方案中采用的模型是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指数序列。

学习目标

学习TOP方案的模型构建方法

内容介绍

数据处理
- 数据扁平化
- 空值填充
- 构造数据集
模型构建
- 构造评估函数
- 模型构造
- 模型训练
- 模型评估
总结

代码示例

数据处理

该TOP方案的数据处理主要包括三部分：

数据扁平化。
空值填充。
构造数据集

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 inlineimport 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 ReduceLROnPlateaufrom sklearn.metrics import mean_squared_error

# 固定随机种子
SEED = 22def 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 = Trueseed_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')

数据扁平化

采用滑窗构造数据集。该方案中只使用了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).datablocks.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

构造数据集

构造训练和验证集。注意这里取每条输入数据的序列长度是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)

模型构建

# 读取数据集
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 = 2trainset = 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)

构造评估函数

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# RMSErmse_scores = 0a = [1.5] * 4 + [2] * 7 + [3] * 7 + [4] * 6y_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 / fenmuaccskill_score += a[i] * np.log(i+1) * cor_irmse_score = rmse(y_true[:, i], y_preds[:, i])rmse_scores += rmse_scorereturn 2/3.0 * accskill_score - rmse_scores

模型构造

不同于前两个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

ConvLSTM模型

LSTM模型是非常经典的时序模型，三个门的结构使得它在挖掘长期的时间依赖任务中有不俗的表现，并且相较于RNN，LSTM能够有效地避免梯度消失问题。对于单个输入样本，在每个时间步上，LSTM的每个门实际是对输入向量做了一个全连接，那么对应到我们这个赛题上，输入X的形状是（N,T,H,W,C），则单个输入样本在每个时间步上输入LSTM的就是形状为（H,W,C）的空间信息。我们知道，全连接网络对于这种空间信息的提取能力并不强，转换成卷积操作后能够在大大减少参数量的同时通过堆叠多层网络逐步提取出更复杂的特征，到这里就可以很自然地想到，把LSTM中的全连接操作转换为卷积操作，就能够适用于时空序列问题。ConvLSTM模型就是这么做的，实践也表明这样的作法是非常有效的。

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所获得的隐藏层状态HtH_tHt和上一个时间步的记忆Mt−1M_{t-1}Mt−1作为输入，首先将HtH_tHt通过自注意力机制得到特征ZhZ_hZh，自注意力机制能够增加HtH_tHt中与其他部分更相关的部分的权重，同时HtH_tHt也作为Query与Mt−1M_{t-1}Mt−1共同通过注意力机制得到特征ZmZ_mZm，用以增强对Mt−1M_{t-1}Mt−1中与HtH_tHt有更强依赖关系的部分的权重，将ZhZ_hZh和ZmZ_mZm拼接起来就得到了二者的聚合特征ZZZ。此时，聚合特征ZZZ中既包含了当前时间步的信息，又包含了全局的时空记忆信息，接下来借鉴LSTM中的门控结构用聚合特征ZZZ对隐藏层状态和记忆单元进行更新，就得到了更新后的隐藏层状态Ht^\hat{H_t}Ht^和当前时间步的记忆MtM_tMt。SAM模块的公式如下：

it′=σ(Wm;zi∗Z+Wm;hi∗Ht+bm;i)gt′=tanh(Wm;zg∗Z+Wm;hg∗Ht+bm;g)Mt=(1−it′)∘Mt−1+it′∘gt′ot′=σ(Wm;zo∗Z+Wm;ho∗Ht+bm;o)Ht^=ot′∘Mt\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} it′=σ(Wm;zi∗Z+Wm;hi∗Ht+bm;i)gt′=tanh(Wm;zg∗Z+Wm;hg∗Ht+bm;g)Mt=(1−it′)∘Mt−1+it′∘gt′ot′=σ(Wm;zo∗Z+Wm;ho∗Ht+bm;o)Ht^=ot′∘Mt

关于注意力机制和自注意力机制可以参考以下链接：

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

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/根号Cscores = 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] // 2self.d_model = d_modelself.input_dim = input_dim# 用1*1卷积实现全连接操作WhHtself.conv_h = nn.Conv2d(input_dim, d_model*3, kernel_size=1)# 用1*1卷积实现全连接操作WmMt-1self.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_dimself.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、Vhhq, hk, hv = torch.split(self.conv_h(h), self.d_model, dim=1)# 同样的方法得到Km和Vmmk, mv = torch.split(self.conv_m(m), self.d_model, dim=1)N, C, H, W = hq.size()# 通过自注意力机制得到ZhZh = attn(hq.view(N, C, -1), hk.view(N, C, -1), hv.view(N, C, -1))  # (N, C, S), C=d_model# 通过注意力机制得到ZmZm = attn(hq.view(N, C, -1), mk.view(N, C, -1), mv.view(N, C, -1))  # (N, C, S), C=d_model# 将Zh和Zm拼接起来，并进行全连接操作得到聚合特征ZZ = 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'ti, g, o = torch.split(self.conv_output(torch.cat([Z, h], dim=1)), self.input_dim, dim=1)  # (N, C, H, W), C=input_dimi = torch.sigmoid(i)g = torch.tanh(g)# 得到更新后的记忆单元Mtm_next = i * g + (1 - i) * m# 得到更新后的隐藏状态Hth_next = torch.sigmoid(o) * m_nextreturn 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_dimself.hidden_dim = hidden_dimpad = kernel_size[0] // 2, kernel_size[1] // 2# 卷积操作Wx*Xt+Wh*Ht-1self.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.deviceN, _, H, W = inputs.size()# 初始化隐藏层状态Htself.hidden_state = torch.zeros(N, self.hidden_dim, H, W, device=device)# 初始化记忆细胞状态ctself.cell_state = torch.zeros(N, self.hidden_dim, H, W, device=device)# 初始化记忆单元状态Mtself.memory_state = torch.zeros(N, self.hidden_dim, H, W, device=device)def forward(self, inputs, first_step=False):# 如果当前是第一个时间步，初始化Ht、ct、Mtif first_step:self.initialize(inputs)# ConvLSTM部分# 拼接Xt和Htcombined = 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、gtcc_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·gtself.cell_state = f * self.cell_state + i * g# 得到当前时间步的隐藏层状态Ht=ot·tanh(ct)self.hidden_state = o * torch.tanh(self.cell_state)# SAM部分，更新Ht和Mtself.hidden_state, self.memory_state = self.sa(self.hidden_state, self.memory_state)return self.hidden_state

在Seq2Seq模型的训练中，有两种训练模式。一是Free running，也就是传统的训练方式，以上一个时间步的输出yt−1^\hat{y_{t-1}}yt−1^作为下一个时间步的输入，但是这种做法存在的问题是在训练的初期所得到的yt−1^\hat{y_{t-1}}yt−1^与实际标签yt−1y_{t-1}yt−1相差甚远，以此作为输入会导致后续的输出越来越偏离我们期望的预测标签。于是就产生了第二种训练模式——Teacher forcing。

Teacher forcing就是直接使用实际标签yt−1y_{t-1}yt−1作为下一个时间步的输入，由老师（ground truth）带领着防止模型越走越偏。但是老师不能总是手把手领着学生走，要逐渐放手让学生自主学习，于是我们使用Scheduled Sampling来控制使用实际标签的概率。我们用ratio来表示Scheduled Sampling的比例，在训练初期，ratio=1，模型完全由老师带领着，随着训练论述的增加，ratio以一定的方式衰减（该方案中使用线性衰减，ratio每次减小一个衰减率decay_rate），每个时间步以ratio的概率从伯努利分布中提取二进制随机数0或1，为1时输入就是实际标签yt−1y_{t-1}yt−1，否则输入为yt−1^\hat{y_{t-1}}yt−1^。

# 构建SA-ConvLSTM模型
class SAConvLSTM(nn.Module):def __init__(self, input_dim, hidden_dim, d_attn, kernel_size):super().__init__()self.input_dim = input_dimself.hidden_dim = hidden_dimself.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 forcingif 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 = Falseelse:teacher_forcing = Falsetotal_steps = input_frames + future_frames - 1outputs = [None] * total_steps# 对于每一个时间步for t in range(total_steps):# 在前12个月，使用每个月的输入样本Xtif 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_frameselse: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)

SAConvLSTM((layers): ModuleList((0): SAConvLSTMCell((conv): Conv2d(65, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))(sa): SAAttnMem((conv_h): Conv2d(64, 96, kernel_size=(1, 1), stride=(1, 1))(conv_m): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1))(conv_z): Conv2d(64, 32, kernel_size=(1, 1), stride=(1, 1))(conv_output): Conv2d(96, 192, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))))(1): SAConvLSTMCell((conv): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))(sa): SAAttnMem((conv_h): Conv2d(64, 96, kernel_size=(1, 1), stride=(1, 1))(conv_m): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1))(conv_z): Conv2d(64, 32, kernel_size=(1, 1), stride=(1, 1))(conv_output): Conv2d(96, 192, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))))(2): SAConvLSTMCell((conv): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))(sa): SAAttnMem((conv_h): Conv2d(64, 96, kernel_size=(1, 1), stride=(1, 1))(conv_m): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1))(conv_z): Conv2d(64, 32, kernel_size=(1, 1), stride=(1, 1))(conv_output): Conv2d(96, 192, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))))(3): SAConvLSTMCell((conv): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))(sa): SAAttnMem((conv_h): Conv2d(64, 96, kernel_size=(1, 1), stride=(1, 1))(conv_m): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1))(conv_z): Conv2d(64, 32, kernel_size=(1, 1), stride=(1, 1))(conv_output): Conv2d(96, 192, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)))))(conv_output): Conv2d(64, 1, kernel_size=(1, 1), stride=(1, 1))
)

模型训练

# 采用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 = 0for 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 = 0with torch.no_grad():for i, data in tqdm(enumerate(validloader)):data, labels = datadata = 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 = scheckpoint = {'best_score': s,'state_dict': model.state_dict()}torch.save(checkpoint, model_weights)

Epoch: 1/5100%|██████████| 1600/1600 [21:43<00:00,  1.23it/s]Training Loss: 3.28950it [00:11,  4.47it/s]Validation Loss: 44.009
Score: -43.458
Epoch: 2/5100%|██████████| 1600/1600 [21:43<00:00,  1.23it/s]Training Loss: 3.08450it [00:11,  4.33it/s]Validation Loss: 25.011
Score: -19.966
Epoch: 3/5100%|██████████| 1600/1600 [21:46<00:00,  1.22it/s]Training Loss: 13.46150it [00:12,  4.16it/s]Validation Loss: 15.438
Score: -14.139
Epoch: 4/5100%|██████████| 1600/1600 [21:54<00:00,  1.22it/s]Training Loss: 17.62750it [00:12,  3.99it/s]Validation Loss: 15.389
Score: -22.500
Epoch: 5/5100%|██████████| 1600/1600 [21:55<00:00,  1.22it/s]Training Loss: 17.59250it [00:11,  4.48it/s]Validation Loss: 15.252
Score: -14.459

# 绘制训练/验证曲线
def training_vis(train_losses, valid_losses):# 绘制损失函数曲线fig = plt.figure(figsize=(8,4))# subplot lossax1 = 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)

模型评估

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

# 加载得分最高的模型
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 = datadata = 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))

总结

这一次的TOP方案没有自己设计模型，而是使用了目前时空序列预测领域现有的模型，另一组TOP选手“ailab”也使用了现有的模型PredRNN++，关于时空序列预测领域的一些比较经典的模型可以参考https://www.zhihu.com/column/c_1208033701705162752

作业

该TOP方案中以sst作为预测目标，间接计算nino3.4指数，学有余力的同学可以尝试用SA-ConvLSTM模型直接预测nino3.4指数。

参考文献

吴先生的队伍方案分享：https://tianchi.aliyun.com/forum/postDetail?spm=5176.12586969.1002.9.561d5330dF9lX1&postId=231465
ailab团队思路分享：https://tianchi.aliyun.com/forum/postDetail?spm=5176.12586969.1002.15.561d5330dF9lX1&postId=210734

【算法竞赛学习】气象海洋预测-Task5 模型建立之 SA-ConvLSTM相关推荐

【算法竞赛学习】气象海洋预测-Task4 模型建立之 TCNN+RNN
气象海洋预测-Task4 模型建立之 TCNN+RNN 该方案中采用的模型是TCNN+RNN. 在Task3中我们学习了CNN+LSTM模型,但是LSTM层的参数量较大,这就带来以下问题:一是参数量大 ...
【算法竞赛学习】数字中国创新大赛智慧海洋建设-Task5模型融合
智慧海洋建设-Task5 模型融合 5.1 学习目标学习融合策略完成相应学习打卡任务 5.2 内容介绍 https://mlwave.com/kaggle-ensembling-guide/ ht ...
【算法竞赛学习】气象海洋预测-Task2 数据分析
气象海洋预测-Task2 数据分析数据分析是我们解决一个数据挖掘任务的重要一环,通过数据分析,我们可以了解标签的分布.数据中存在的缺失值和异常值.特征与标签之间的相关性.特征之间的相关性等,并根据数 ...
【算法竞赛学习】二手车交易价格预测-Task5模型融合
二手车交易价格预测-Task5 模型融合五.模型融合 Tip:此部分为零基础入门数据挖掘的 Task5 模型融合部分,带你来了解各种模型结果的融合方式,在比赛的攻坚时刻冲刺Top,欢迎大家后续多多 ...
【算法竞赛学习】气象海洋预测-Task1 气象数据分析常用工具
气象海洋预测-Task1 气象数据分析常用工具气象科学中的数据通常包含多个维度,例如本赛题中给出的数据就包含年.月.经度.纬度四个维度,为了便于数据的读取和操作,气象数据通常采用netCDF文件来存 ...
【算法竞赛学习】数字中国创新大赛智慧海洋建设-Task4模型建立
智慧海洋建设-Task4模型建立此部分为智慧海洋建设竞赛的模型建立模块.在该模块中主要介绍了如何进行模型建立并对模型调优. 学习目标学习如何选择合适的模型以及如何通过模型来进行特征选择掌握随机森 ...
Kaggle泰坦尼克号生存预测挑战——模型建立、模型调参、融合
Kaggle泰坦尼克号生存预测挑战这是kaggle上Getting Started 的Prediction Competition,也是比较入门和简单的新人赛,我的最好成绩好像有进入top8%,重新 ...
【算法竞赛学习】金融风控之贷款违约预测-模型融合
Task5 模型融合 Tip:此部分为零基础入门金融风控的 Task5 模型融合部分,欢迎大家后续多多交流. 赛题:零基础入门数据挖掘 - 零基础入门金融风控之贷款违约预测项目地址:https:// ...
【算法竞赛学习】数字中国创新大赛智慧海洋建设-Task3特征工程
智慧海洋建设-Task3 特征工程此部分为智慧海洋建设竞赛的特征工程模块,通过特征工程,可以最大限度地从原始数据中提取特征以供算法和模型使用.通俗而言,就是通过X,创造新的X'以获得更好的训练.预测 ...

【算法竞赛学习】气象海洋预测-Task5 模型建立之 SA-ConvLSTM