“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 方案的数据处理主要包括三部分:
数据扁平化。
空值填充。
构造数据集
项目链接以及码源见文末
在该方案中除了没有构造新的特征外,其他数据处理方法都与 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 模型全称时间卷积网络(Temporal Convolutional Network),与 RNN 一样是时序模型。TCN 以 CNN 为基础,为了适应序列问题,它从以下三方面做出了改进:
因果卷积
TCN 处理输入与输出等长的序列问题,它的每一个隐藏层节点数与输入步长是相同的,并且隐藏层 t 时刻节点的值只依赖于前一层 t 时刻及之前节点的值。也就是说 TCN 通过追溯前因(t 时刻及之前的值)来获得当前结果,称为因果卷积。
扩张卷积
传统 CNN 的感受野受限于卷积核的大小,需要通过增加池化层来获得更大的感受野,但是池化的操作会带来信息的损失。为了解决这个问题,TCN 采用扩张卷积来增大感受野,获取更长时间的信息。扩张卷积对输入进行间隔采样,采样间隔由扩张因子 d 控制,公式定义如下:
F(s)=(X∗df)(s)=i=0∑k−1f(i)×Xs−di
其中 X 为当前层的输入,k 为当前层的卷积核大小,s 为当前节点的时刻。也就是说,对于扩张因子为 d、卷积核为 k 的隐藏层,对前一层的输入每 d 个点采样一次,共采样 k 个点作为当前时刻 s 的输入。这样 TCN 的感受野就由卷积核的大小 k 和扩张因子 d 共同决定,可以获取更长时间的依赖信息。
残差连接
网络的层数越多,所能提取到的特征就越丰富,但这也会带来梯度消失或爆炸的问题,目前解决这个问题的一个有效方法就是残差连接。TCN 的残差模块包含两层卷积操作,并且采用了 WeightNorm 和 Dropout 进行正则化,如下图所示。
总的来说,TCN 是卷积操作在序列问题上的改进,具有 CNN 参数量少的优点,可以搭建更深层的网络,相比于 RNN 不容易存在梯度消失和爆炸的问题,同时 TCN 具有灵活的感受野,能够适应不同的任务,在许多数据集上的比较表明 TCN 比 RNN、LSTM、GRU 等序列模型有更好的表现。
想要更深入地了解 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 单元结构与 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
复制代码
将 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 部分
TCNN 部分的模型结构类似传统 CNN 的结构,非常规整,通过逐渐增加通道数来提取更丰富的特征表达。需要注意的是输入数据的格式是(N,T,H,W,C),为了匹配卷积层的输入格式,需要将数据格式转换为(N,T,C,H,W)。
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.adaptive_avg_pool2d.html?highlight=adaptive_avg_pool2d#torch.nn.functional.adaptive_avg_pool2d
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 进行信息提取,综合三种不同时间尺度的序列信息得到最终的预测序列。
该方案同样选择了自己设计模型,在构造模型时充分考虑了数据集情况和问题背景,并能灵活运用各种网络层来处理特定问题,这种模型构造思路要求对不同网络层的作用有较为深刻地理解,方案中各种网络层的用法值得大家学习和借鉴。
参考文献
Top1 思路分享:https://tianchi.aliyun.com/forum/postDetail?spm=5176.12586969.1002.6.561d482cp7CFlx&postId=210391
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
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 所获得的隐藏层状态Ht和上一个时间步的记忆Mt−1作为输入,首先将Ht通过自注意力机制得到特征Zh,自注意力机制能够增加Ht中与其他部分更相关的部分的权重,同时Ht也作为 Query 与Mt−1共同通过注意力机制得到特征Zm,用以增强对Mt−1中与Ht有更强依赖关系的部分的权重,将Zh和Zm拼接起来就得到了二者的聚合特征Z。此时,聚合特征Z中既包含了当前时间步的信息,又包含了全局的时空记忆信息,接下来借鉴 LSTM 中的门控结构用聚合特征Z对隐藏层状态和记忆单元进行更新,就得到了更新后的隐藏层状态Ht^和当前时间步的记忆Mt。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}$$
关于注意力机制和自注意力机制可以参考以下链接:
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,也就是传统的训练方式,以上一个时间步的输出yt−1^作为下一个时间步的输入,但是这种做法存在的问题是在训练的初期所得到的yt−1^与实际标签yt−1相差甚远,以此作为输入会导致后续的输出越来越偏离我们期望的预测标签。于是就产生了第二种训练模式——Teacher forcing。
Teacher forcing 就是直接使用实际标签yt−1作为下一个时间步的输入,由老师(ground truth)带领着防止模型越走越偏。但是老师不能总是手把手领着学生走,要逐渐放手让学生自主学习,于是我们使用 Scheduled Sampling 来控制使用实际标签的概率。我们用 ratio 来表示 Scheduled Sampling 的比例,在训练初期,ratio=1,模型完全由老师带领着,随着训练论述的增加,ratio 以一定的方式衰减(该方案中使用线性衰减,ratio 每次减小一个衰减率 decay_rate),每个时间步以 ratio 的概率从伯努利分布中提取二进制随机数 0 或 1,为 1 时输入就是实际标签yt−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_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
项目链接以及码源见文末
链接跳转到文末,欢迎订阅
更多文章请关注公重号:汀丶人工智能
评论