【Datawhale AI 夏令营】AI极端降水预报赛道个人笔记

血小板3512

于 2024-08-03 23:56:13 发布

阅读量297

点赞数 4

分类专栏：时空数据处理文章标签：人工智能 pytorch

本文链接：https://blog.csdn.net/m0_56172605/article/details/140898855

版权

时空数据处理专栏收录该内容

1 篇文章 0 订阅

订阅专栏

环境：

1.Windows环境：Windows11
2.运行环境：python-3.9.10
numpy-1.23.0
torch-1.10.0+cu102

一、ConvLSTM模型

1.为何用它

第一次看到这份baseline的时候，我很诧异，原来一个简单的卷积层就能做降水预报这样子高难度的事情。通过加深模型以及增多epoch轮数，比赛成绩切实有效地蹭蹭往上涨。然后，我开始想，能不能使用其他模型来做这个任务，而第一次直播课告诉了我这个答案——ConvLSTM模型。

2.模型简介

我上网看到了一篇博客，它让我知道了ConvLSTM模型提出的目的就是为了解决降水预报问题，同时，模型还可以解决其他时空序列的预测问题。ConvLSTM模型不仅像LSTM一样可以捕获时间特征，还可以像CNN一样捕获空间特征，并且作者通过实验证明了ConvLSTM在获取时空关系上比LSTM有更好的效果。

3.模型搭建

我仔细阅读了上面博客的代码，并结合chatGPT给出的示例代码，写了如下代码：

（1）ConvLSTMCell

import torch  
import torch.nn as nn

class ConvLSTMCell(nn.Module):
    def __init__(self, input_dim, hidden_dim, kernel_size, bias=True):
        super(ConvLSTMCell, self).__init__()
        self.input_dim = input_dim
        self.hidden_dim = hidden_dim  
        self.kernel_size = kernel_size
        self.bias = bias
        self.conv = nn.Conv2d(in_channels=input_dim+hidden_dim,
                              out_channels=4*hidden_dim,
                              kernel_size=kernel_size,
                              padding=kernel_size//2,
                              bias=bias)

    def forward(self, x, h, c):  
        combined = torch.cat((x, h), dim=1)
        combined_conv = self.conv(combined)
        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)     # 单元状态
        
        c_next = f * c + i * g
        h_next = o * torch.tanh(c_next)  
        return h_next, c_next

（2）ConvLSTM

class ConvLSTM(nn.Module):  
    def __init__(self, input_dim, hidden_dim, kernel_size, num_layers, bias=True):  
        super(ConvLSTM, self).__init__()
        self.layers = nn.ModuleList()  

        for i in range(num_layers):
            input_size = input_dim if i == 0 else hidden_dim
            cell = ConvLSTMCell(input_dim=input_size, hidden_dim=hidden_dim, kernel_size=kernel_size, bias=bias)
            self.layers.append(cell)
    
    def forward(self, x):
        # Get size
        batch_size, time_steps, channels, height, width = x.size()
        # Init hidden
        h, c = [None] * len(self.layers), [None] * len(self.layers)
        for i, layer in enumerate(self.layers):
            h[i], c[i] = (torch.zeros(batch_size, layer.hidden_dim, height, width).to(x.device),  
                          torch.zeros(batch_size, layer.hidden_dim, height, width).to(x.device))
        # Init input data
        layer_output_list = []
        current_layer_input = x
        for i, layer in enumerate(self.layers):  
            current_layer_output = []
            for t in range(time_steps):
                h[i], c[i] = layer(current_layer_input[:, t, :, :, :], h[i], c[i])
                current_layer_output.append(h[i])
            # Output data → Next input data
            current_layer_input = torch.stack(current_layer_output, dim=1)
            layer_output_list.append(current_layer_input)
        # Return last layer output
        return layer_output_list[-1]

（3）ConvLSTMModel

class ConvLSTMModel(nn.Module):  
    def __init__(self, input_dim, hidden_dim, kernel_size, num_layers, time_steps, output_dim):  
        super(ConvLSTMModel, self).__init__()
        self.convlstm = ConvLSTM(input_dim, hidden_dim, kernel_size, num_layers)  
        self.conv = nn.Conv2d(time_steps*hidden_dim, time_steps*output_dim, 3, 1, 1)

    def forward(self, x):  
        last_output = self.convlstm(x)
        batch_size, time_steps, hidden_dim, height, width = tuple(last_output.shape)
        last_output = last_output.reshape(batch_size, -1, height, width)
        output = self.conv(last_output)
        output = output.reshape(batch_size, time_steps, -1, height, width)
        return output

（4）训练模型

# 使用示例
input_shape = (10, 5, 3, 64, 64)
X_train = torch.rand(input_shape).float()
y_train = torch.randn(10, 5, 1, 64, 64)

# 创建模型
model = ConvLSTMModel(input_dim=3, hidden_dim=16, kernel_size=3, num_layers=2, time_steps=5, output_dim=1)
criterion = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

# 训练模型
model.train()
num_epochs = 100
for epoch in range(num_epochs):  
    optimizer.zero_grad()
    outputs = model(X_train)
    loss = criterion(outputs, y_train)  
    loss.backward()  
    optimizer.step()  
    print(f'Epoch [{epoch + 1}/{num_epochs}], Loss: {loss.item():.4f}')

二、经历感悟

通过参加比赛，让我真真切切学习到了课本之外的内容，听大佬的解说，更加深入地理解五个上分思路（数据、模型、损失函数、训练方式、超参数）的做法。
对于数据，可以采用更多的数据，以及不使用全部的特征。
对于模型，可以思考如何加深模型，但是加深模型并不一定会带来效果提升。另外，比赛禁止预训练模型，但可以使用预训练模型结构，再随机初始化参数。
对于损失函数，并非只能使用传统的MSE，还可以使用SmoothL1Loss()替代。另外，赛事评测使用的是CSI指标，可以自己编写损失函数去接近原有的任务。
对于训练方式，可以将所有数据都作为训练数据，也可以划分训练集和验证集，然后选择验证损失最小的模型。