DataWhaleAI夏令营：极端降水预测baseline

本文主要介绍一下baseline以及运行baseline过程中遇到的问题和解决方法，还有后续的提分策略。

1.环境配置

在魔搭社区官网搭建的环境中能很快的完成以下库的安装。

!pip install xarray[complete]
!apt update&&apt install axel

考虑到该网站不能保存环境和数据，便打算本地搭建环境，但是安装xarray[complete]包时遇到以下报错（目前还未解决）：

File "/data/miniconda/envs/torch/lib/python3.10/site-packages/pip/_vendor/urllib3/response.py", line 443, in _error_catcher
    raise ReadTimeoutError(self._pool, None, "Read timed out.")
pip._vendor.urllib3.exceptions.ReadTimeoutError: HTTPSConnectionPool(host='files.pythonhosted.org', port=443): Read timed out.

2.数据集下载与解压（以2021年为例）

!axel -n 5 -o weather.round1.train.gt.2019-2021.zip 'https://tianchi-race-prod-sh.oss-cn-shanghai.aliyuncs.com/file/race/documents/532234/bigFile/weather.round1.train.gt.2019-2021.zip?Expires=1722203212&OSSAccessKeyId=LTAI5t7fj2oKqzKgLGz6kGQc&Signature=xz9ExPLpJRyCYzNYJVO3dTgu3AM%3D&response-content-disposition=attachment%3B%20'
!axel -n 5 -o weather.round1.train.ft.2021.1.zip 'https://tianchi-race-prod-sh.oss-cn-shanghai.aliyuncs.com/file/race/documents/532234/bigFile/weather.round1.train.ft.2021.1.zip?Expires=1722203232&OSSAccessKeyId=LTAI5t7fj2oKqzKgLGz6kGQc&Signature=hCntmc6xHNL%2FqyqeRZGiK8x9JyA%3D&response-content-disposition=attachment%3B%20'
!axel -n 5 -o weather.round1.train.ft.2021.2.zip 'https://tianchi-race-prod-sh.oss-cn-shanghai.aliyuncs.com/file/race/documents/532234/bigFile/weather.round1.train.ft.2021.2.zip?Expires=1722203248&OSSAccessKeyId=LTAI5t7fj2oKqzKgLGz6kGQc&Signature=NF8BJgoVLhIOsPXFyKu6zPIl8TU%3D&response-content-disposition=attachment%3B%20'
!axel -n 5 -o weather.round1.test.zip 'https://tianchi-race-prod-sh.oss-cn-shanghai.aliyuncs.com/file/race/documents/532234/bigFile/weather.round1.test.zip?Expires=1722203263&OSSAccessKeyId=LTAI5t7fj2oKqzKgLGz6kGQc&Signature=JmtcVll3Pk4bfsxiWFhw3HqGVrc%3D&response-content-disposition=attachment%3B%20'

!unzip -q -n weather.round1.train.ft.2021.1.zip -d feature
!unzip -q -n weather.round1.train.ft.2021.2.zip -d feature

2021年的数据下载后发现分上半年和下半年，考虑到后续代码运行读取路径问题，我们要将其放置在同一文件夹下，可在终端feature目录下依次执行以下代码即可完成：

mv ./2021.1/* ./2021
mv ./2021.2/* ./2021

3.数据集路径配置

# path config
feature_path = 'feature'
gt_path = 'groundtruth'
years = ['2021']
fcst_steps = list(range(1, 73, 1))

• 比赛的数据部分分为数据特征(feature)和数据真值(groundtruth)两部分，数据特征是模型训练的输入，数据真值是模型训练的标签

• 其中数据特征部分输入的路径目录下包含年份文件夹

• 数据真值部分输入的路径目录下包含3个年份的.nc数据，以2021年为例，就在years中添加2021

• fcst_steps指预测的时间步长, 从第1小时到第72小时, 间隔为1小时

4.定义Featue类和GroundTruth类

# Feature部分
class Feature:
    def __init__(self):
        self.path = feature_path
        self.years = years
        self.fcst_steps = fcst_steps
        self.features_paths_dict = self.get_features_paths()

    def get_features_paths(self):
        init_time_path_dict = {}
        for year in self.years:
            init_time_dir_year = os.listdir(os.path.join(self.path, year))
            for init_time in sorted(init_time_dir_year):
                init_time_path_dict[pd.to_datetime(init_time)] = os.path.join(self.path, year, init_time)
        return init_time_path_dict

    def get_fts(self, init_time):
        return xr.open_mfdataset(self.features_paths_dict.get(init_time) + '/*').sel(lead_time=self.fcst_steps).isel(
            time=0)

 Feature类用于处理特征数据，这些数据通常作为模型的输入。

• __init__方法：初始化Feature实例，接收特征数据的路径(feature_path)，年份列表(years)，和预测步长(fcst_steps)。然后调用get_features_paths方法来构建一个字典，该字典以初始化时间作为键，特征数据路径作为值。

• get_features_paths方法：遍历指定路径下的所有年份，然后列出每个年份下的所有初始化时间的目录。对于每个初始化时间，将对应的数据路径添加到字典中。

• get_fts方法：根据给定的初始化时间(init_time)，打开对应的特征数据集，并选择特定的预测步长(fcst_steps)和时间点（通常是0，表示初始化时间点）。

# GroundTruth部分
class GT:
    def __init__(self):
        self.path = gt_path
        self.years = years
        self.fcst_steps = fcst_steps
        self.gt_paths = [os.path.join(self.path, f'{year}.nc') for year in self.years]
        self.gts = xr.open_mfdataset(self.gt_paths)

    def parser_gt_timestamps(self, init_time):
        return [init_time + pd.Timedelta(f'{fcst_step}h') for fcst_step in self.fcst_steps]

    def get_gts(self, init_time):

        return self.gts.sel(time=self.parser_gt_timestamps(init_time))

 GT类用于处理真实标签数据，这些数据通常用于模型训练过程中的损失计算和模型评估。

• __init__方法：初始化GT实例，接收真实标签数据的路径(gt_path)，年份列表(years)，和预测步长(fcst_steps)。然后构建一个包含所有年份真实标签文件路径的列表，并使用xarray的open_mfdataset方法打开这些数据集。

• parser_gt_timestamps方法：根据初始化时间和预测步长，解析出所有预测时间点。

• get_gts方法：根据给定的初始化时间(init_time)，从真实标签数据集中选择对应的时间点。

5.定义mydataset类

# 构建Dataset部分
class mydataset(Dataset):
    def __init__(self):
        self.ft = Feature()
        self.gt = GT()
        self.features_paths_dict = self.ft.features_paths_dict
        self.init_times = list(self.features_paths_dict.keys())

    def __getitem__(self, index):
        init_time = self.init_times[index]
        ft_item = self.ft.get_fts(init_time).to_array().isel(variable=0).values
        # print(type(ft_item))
        gt_item = self.gt.get_gts(init_time).to_array().isel(variable=0).values
        # print(type(gt_item))
        return ft_item, gt_item

    def __len__(self):
        return len(list(self.init_times))

• class mydataset(Dataset)：定义了一个名为 mydataset 的类，它继承自 PyTorch 的 Dataset 类。这使得 mydataset 可以被用作 PyTorch 的数据加载器（如 DataLoader）。

 __init__ 方法

• self.ft = Feature()：创建一个 Feature 类的实例，用于处理特征数据。
• self.gt = GT()：创建一个 GT 类的实例，用于处理真实标签数据。
• self.features_paths_dict = self.ft.features_paths_dict：获取特征数据路径字典，该字典由 Feature 类的 get_features_paths 方法生成。
• self.init_times = list(self.features_paths_dict.keys())：获取所有初始化时间点的列表，这些时间点是特征数据路径字典的键。

 __getitem__ 方法

• def __getitem__(self, index)：这个方法定义了如何获取数据集中的一个项。它接收一个索引 index，然后返回对应的特征数据和真实标签数据。
• init_time = self.init_times[index]：根据索引获取对应的初始化时间点。
• ft_item = self.ft.get_fts(init_time).to_array().isel(variable=0).values：调用 Feature 类的 get_fts 方法获取特征数据，然后将其转换为数组，并选择第一个变量（如果有多个变量的话），最后获取其值。
• gt_item = self.gt.get_gts(init_time).to_array().isel(variable=0).values：调用 GT 类的 get_gts 方法获取真实标签数据，然后同样转换为数组，并选择第一个变量，获取其值。
• return ft_item, gt_item：返回一个包含特征数据和真实标签数据的元组。

 __len__ 方法

• def __len__(self)：这个方法定义了数据集的长度，即数据集中有多少个初始化时间点。
• return len(list(self.init_times))：返回初始化时间点列表的长度。

6.构建模型

以下模型作为一个基础的入门模型，仅定义了一层卷积层

# 模型构建部分
import torch.nn as nn

class Model(nn.Module):
    def __init__(self, num_in_ch, num_out_ch):
        super(Model, self).__init__()
        self.conv1 = nn.Conv2d(num_in_ch, num_out_ch, 3, 1, 1)

    def forward(self, x):
        B, S, C, W, H = tuple(x.shape)
        x = x.reshape(B, -1, W, H)
        out = self.conv1(x)
        out = out.reshape(B, S, W, H)
        return out

# define model
in_varibales = 24  # 输入变量的数量
in_times = len(fcst_steps)  # 输入时间步长的数量
out_varibales = 1  # 输出变量的数量
out_times = len(fcst_steps)  # 输出时间步长的数量
input_size = in_times * in_varibales  # 输入特征的总维度
output_size = out_times * out_varibales  # 输出特征的总维度
model = Model(input_size, output_size).cuda()

下面定义了模型的损失函数 ，用于模型训练做反向传播：

# define loss
loss_func = nn.MSELoss()

7.模型训练

import numpy as np
import torch
# from tqdm import tqdm
# Train the model
num_epochs = 1  # 设置训练的轮数
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)  # 使用Adam优化器

# for epoch in tqdm(range(num_epochs)):
for epoch in range(num_epochs):
    for index, (ft_item, gt_item) in enumerate(train_loader):
        ft_item = ft_item.cuda().float()  # 将特征数据移动到GPU并转换为float类型
        gt_item = gt_item.cuda().float()  # 将目标数据移动到GPU并转换为float类型
        
        # Forward pass
        output_item = model(ft_item)  # 通过模型进行前向传播
        loss = loss_func(output_item, gt_item)  # 计算损失函数
        
        # Backward and optimize
        optimizer.zero_grad()  # 清除之前的梯度
        loss.backward()  # 反向传播，计算梯度
        optimizer.step()  # 更新模型的权重
        
        # Print the loss for every 10 steps
        if (index+1) % 10 == 0:
            print(f"Epoch [{epoch+1}/{num_epochs}], Step [{index+1}/{len(train_loader)}], Loss: {loss.item():.4f}")

# Save the model weights
torch.save(model.state_dict(), 'model_weights.pth')

训练过程中会遇到以下报错：

KeyError: "not all values found in index 'time'"

这是因为数据集中存在缺失值，导致的时间序列不连续，此时需要对前面的mydataset类做出调整：

# 构建Dataset部分
class mydataset(Dataset):
    def __init__(self):
        self.ft = Feature()
        self.gt = GT()
        self.features_paths_dict = self.ft.features_paths_dict
        self.init_times = list(self.features_paths_dict.keys())

    def __getitem__(self, index):
        init_time = self.init_times[index]
        try:
            ft_item = self.ft.get_fts(init_time).to_array().isel(variable=0).values
            gt_item = self.gt.get_gts(init_time).to_array().isel(variable=0).values
        except KeyError as e:    
            print(e)
            print(f'init_time: {init_time} not found')
            # return None, None
            return self.__getitem__(index - 1)
        
        return ft_item, gt_item

    def __len__(self):
        return len(list(self.init_times))

其中下面的这一部分有效的解决了这一问题，当遇到不连续的时间序列值时，考虑前一个索引，直到索引为有效值：

        except KeyError as e:    
            print(e)
            print(f'init_time: {init_time} not found')
            # return None, None
            return self.__getitem__(index - 1)

8.模型推理

# Inference
# Load the model weights
model.load_state_dict(torch.load('model_weights.pth'))
model.eval()
import os

test_data_path = "test/weather.round1.test"
os.makedirs("./output", exist_ok=True)
for index, test_data_file in enumerate(os.listdir(test_data_path)):
    test_data = torch.load(os.path.join(test_data_path, test_data_file))
    test_data = test_data.cuda().float()
    
    # Forward pass
    output_item = model(test_data)
    
    # Print the output shape
    print(f"Output shape for sample {test_data_file.split('.')[0]}: {output_item.shape}")
    
    # Save the output
    output_path = f"output/{test_data_file}"
    torch.save(output_item.cpu(), output_path)
    # Load the model weights
    model.load_state_dict(torch.load("model_weights.pth"))

通过加载模型使用测试数据作为输入, 得到预测结果文件夹output，将其压缩为output.zip

!zip -r output.zip output

提交至大赛官网，得到baseline的第一个分数0.0047.

9.后续提分思路

通过分数不难发现baseline仅仅给了一个最最基础的运行结果，还有很大的提升空间，后续打算按照以下思路进行提分：

• 特征工程：官网给出了数据集的基本信息，有24个基本特征，其中最后一个特征tp是需要预测的变量，可以考虑一下其余23个特征与tp的相关性，或者尝试添加新特征。

• 搭建更大的模型： baseline只给出了仅有一层卷积层的模型，这是不够的，可以考虑一些现有的时间序列预测模型，适当做修改
• 超参数调整：baseline给的batchsize=1，num_epoch=1，可以考虑适当增加batchsize和num_epoch
• 增大数据量：当前只考虑了2021年的数据，后续可以加入其余2年的数据