这次夏令营我选择了地球科学赛道：AI极端降水预报。学习并实践进行极端降水预测。
Task2：降水预测Baseline详解
搭建模型并解决问题大致分为以下几个步骤:
-
定义数据集, 建立起训练数据和标签之间的关系；定义数据加载器(DataLoader)， 方便取数据进行训练
-
定义模型, 利用PyTorch搭建网络，根据输入输出数据维度实例化模型
-
定义损失函数, 优化器, 训练周期, 训练模型并保存模型参数
-
模型加载及推理(模型预测)，输入测试数据输出要提交的文件
详解：
1、安装运行baseline所必须的库并解压数据集
!pip install xarray[complete]使用pip指令，并且在pip前加上！使其能在notebook中运行
!unzip -q -n weather.round1.train.gt.2019-2021.zip -d groundtruth
!unzip -q -n weather.round1.test.zip -d test解压数据集
2、导入函数库
import os
import pandas as pd
import xarray as xr
from torch.utils.data import Dataset, DataLoader3、数据集的相关配置
feature:数据特征 模型训练的输入
gt:数据真值 模型训练的标签
# path config
feature_path = 'feature'
gt_path = 'groundtruth'
years = ['2021']
fcst_steps = list(range(1, 73, 1))fcst_steps指预测的时间步长, 从第1小时到第72小时, 间隔为1小时
4、Feature类和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的输入路径、使用到的年份、预测的时间点、获取到具体某个feature数据的路径;
-
get_features_paths函数是构造函数中self.features_paths_dict的具体实现, 通过遍历feature-2021(年份)-20210101-00(具体时间戳)，获取到每个具体时间点存放数据特征的路径; -
最后的
get_fts函数, 通过使用xarray库中加载cdf数据，其中sel()中要增加对应的预测预测间隔,isel()中time要设置为0, 表示从头开始组成数据. 这听起来可能有些绕, 我们直接从数据类型上说话.如下图, 我们可以看到在.sel()后, 数据的维度是(1, 6, 24, 57, 81), 其中的6是因为我们看到的feature中的.nc数据, 是每隔6小时一取的；再经过.isel()方法后, 数据维度就变成了(6, 24, 57, 81), 即通过取time=0处索引的数据, 去除了作为数据输入时, 用不到的时间戳，这样实现了数据集的结构化处理。
GT
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))-
构造函数与Feature部分类似, 其中
gt_path如此设计是因为groundtruth文件夹下就有三个.nc文件, 2021.nc等，相比feature, 多出来的self.gts即添加了年份中所有的gt数据, 之后根据需要再去取对应的部分即可 -
parser_gt_timestamps这个函数是根据输入获取到时间戳, 如果我们查看gt中.nc数据的话会发现gt中包含的时间戳及对应数据要比feature中的数据更全, 因此我们采用从feature中获取到时间戳, 然后使用针对的时间戳取对应的groundtruth数据 -
get_gts中, 跟feature中的类似, 取对应时间戳的数据即可, 同样, 我们也可以打印出这个函数的返回值查看数据, 是一个(72, 57, 81)维度的数据:
5、定义数据集
# 构建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))-
构造函数, 分别实例化了Feature和GT类, 定义了取feature的路径还有feature中含有的时间戳数据, 方便在GT中找到对应时间的数据
-
__get_item__进行取数据, 通过传入的index, 获取到对应的时间戳, 通过此时间戳, 从feature和gt中获取到对应的数据. -
__len__()函数返回数据长度
6、模型的构建
首先, 要定义一个模型类, 他是继承于nn.Module类的, 这个类是pytorch中构建网络时最常用的类之一, 方便了我们做初始化操作, 其中super(Model, self).__init__()就是用来使用父类nn.Module中的初始化方式; 在构造函数中, 我们可以看到baseline只简单的定义了一层卷积, 其中的num_in_ch和num_out_ch分别指输入的通道数和输出的通道数(特征数)
接forward()函数是我们在使用网络时做前向传播的函数, 其底层的实现是利用了__call__()这个内置函数, 方便我们实例化模型以后直接可以通过model(inputs)形式调用forward函数。forward函数中, 首先记录了输入数据feature的shape, 然后将输入数据的时间维度和特征维度叠加到一起，这样再经过卷积层得到输出, 得到输出以后再对输出数据做一下维度上的reshape, 保证跟赛题中要求的输出维度一致
定义模型的输入输出维度, 确保模型可以做实例化, 另外时间维度和特征维度需要按照定义模型中的shape操作来对应相乘, 否则就会有维度上的报错； 根据channel进行实例化以后, 因为我们本次使用魔搭的GPU资源, 就需要使用.cuda()将模型放在gpu上
# 模型构建部分
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()定义模型的损失函数, 一般对于回归任务而言, MSELoss是最常用的
# define loss
loss_func = nn.MSELoss()7、模型训练
定义优化器和训练轮数, 使用了Adam优化器, 设置的学习率为0.001
接下来是一个二层的循环, 外层循环是在整个训练轮数上做的, 每次经过这个循环就表示已经遍历了一遍全局数据集；内层循环是根据dataloader中设置的batchsize, 每次循环取batchsize个数据进行训练和优化
获取到feature和gt数据后, 把数据放到GPU上, 保证和模型.cuda()一致，接下来, 将训练数据(feature)输入到模型中获取输出, 得到这个输出以后, 计算输出和gt之间的损失函数, 以便模型进行后续优化
正向传播以后, 梯度回传反向传播来优化我们的模型, 首先使用optimizer.zero_grad()先清空梯度, 再计算模型输出和loss, 然后使用.backward()进行反向传播， optimizer.step()用于更新模型参数.
最后要将模型参数数据(包括权重weight和偏置bias)进行保存, 以便后续的推理使用
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)
# 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')8、模型推理
首先加载模型权重, 使用model.eval()将模型置于推理模式
接下来, 将测试输入数据的路径进行指定
接着一个循环来遍历测试数据路径下所有的数据, 使用torch.load加载；取到输入数据以后, 使用.cuda()将数据放到GPU上，和model保持一致
接着调用模型, 输入测试数据，得出输出数据
最后, 对输出的数据进行保存
# 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"))
扫一扫
806
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
