Datawhale Al夏令营-AI for Science

第三期DatawhaleAI夏令营-AI-降雨量预测

首先，本次学习的主题是深度学习，与上期机器学习存在部分差异。
查看比赛数据发现：本次比赛数据量非常大，单个文件7G，对于传统预测而言无疑是难上加难。
而深度学习在部署的过程中需要大量的数量进行训练，契合海量数据的研究工作。
以下为代码解析：（思路等明天重写）

1. 概述：
   主要实现了从文件中加载特征和真实标签，构建了一个PyTorch的数据集，并训练了一个卷积神经网络模型。

• 首先，定义了两个类Feature和GT，分别用于加载特征数据和真实标签数据。Feature类通过指定的路径和年份，获取特征数据的路径，并提供了根据起始时间获取特征数据的方法。GT类同样通过指定的路径和年份，获取真实标签数据的路径，并提供了根据起始时间获取真实标签数据的方法。
• 接着，定义了一个mydataset类，继承自PyTorch的Dataset类。该类在初始化时加载特征和真实标签的数据路径，并提供了根据索引获取特征和真实标签的方法。
  - 然后，定义了一个卷积神经网络模型Model，该模型包含一个卷积层。在模型训练部分，首先定义了模型、损失函数和优化器。然后，通过遍历数据集进行模型的训练，每次取出一个样本，进行前向传播、计算损失、反向传播和优化。在训练过程中，每隔10个步骤打印一次损失值。
• 最后，进行了模型的推理。从指定路径加载测试数据，对每个测试样本进行前向传播，打印输出的形状，并将输出保存到指定路径。

2. 代码：

# 导入Python包
import os
import torch.nn as nn
import pandas as pd
import xarray as xr
from torch.utils.data import Dataset, DataLoader
# 定义数据存储的路径
feature_path = './feature/' 
gt_path = './groundtruth/'
# select train_data years 
years = ['2021']
fcst_steps = list(range(1, 73, 1))

# 定义了两个类Feature和GT，分别用于加载特征数据和真实标签数据。Feature类通过指定的路径和年份，获取特征数据的路径，并提供了根据起始时间获取特征数据的方法。GT类同样通过指定的路径和年份，获取真实标签数据的路径，并提供了根据起始时间获取真实标签数据的方法。
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)
    

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))


# 定义了一个mydataset类，继承自PyTorch的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))

my_data = mydataset()
print('sample num:', mydataset().__len__())
train_loader = DataLoader(my_data, batch_size=1, shuffle=True)

# 定义了一个卷积神经网络模型Model（该模型包含一个卷积层，简单模型）。
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

# 根据参数实例化模型
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()
loss_func = nn.MSELoss()

# Baseline中为了节约时间，epoch设置为1。
# epoch：完整地遍历一次整个训练数据集的过程
import numpy as np
import torch

num_epochs = 1
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

# 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()
        gt_item = gt_item.cuda().float()
        print(type(ft_item))
        print(type(gt_item))
        
        # 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')

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


# 根据保存的模型权重进行测试数据检验效果，将输出结果保存到本地
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()
    
    # 前向过程
    output_item = model(test_data)
    
    # 查看输出结构
    print(f"Output shape for sample {test_data_file.split('.')[0]}: {output_item.shape}")
    
    # 保存输出
    output_path = f"output/{test_data_file}"
    torch.save(output_item.cpu(), output_path)
    # 加载模型权重
    model.load_state_dict(torch.load("model_weights.pth"))