用Python从头到尾复现一篇Nature的工作 ：
2.神经网络的构建与训练
作者：Vector
前言
本篇文章将从数据下载、处理、神经网络训练、画图四个大步骤叙说笔者在复现 Deep learning for multi-year ENSO forecasts这篇文章的工作。所涉及Python库有 wget , matplotlib , numpy ,xarray , pytorch 等一系列在深度学习以及气象数据处理中经常使用的函数库，希望这篇文章能够对大家有所帮助。笔者也只是大学二年级的本科生，做这些东西也只是凭借个人兴趣，水平低下、错误频出也是常有的事情，请大家见谅。
简单介绍一下这篇文章，这篇文章主要是用 sea surface temp 和 sea surface height 来预测 Nino3.4区的海温（一个厄尔尼诺指数）。此文使用的神经网络、数据的处理都不是很复杂，适合作为气象神经网络入门的第一个尝试性工作。
本文是复现工作的第二篇文章，主要讲解 神经网络构建与训练。
文献地址：Deep learning for multi-year ENSO forecasts | Nature https://www.nature.com/articles/s41586-019-1559-7
上一篇文章地址：用Python复现一篇Nature的研究: 1.数据下载及预处理 https://mp.weixin.qq.com/s/obWDJ6xJrpOL93ep0nG-FA
本工作的所有代码在Github上已经上传lankkk/-Deep-learning-for-multi-year-ENSO-forecasts-: Deep learning for multi-year ENSO forecasts 复现工作相关代码、数据 (github.com) https://github.com/lankkk/-Deep-learning-for-multi-year-ENSO-forecasts-
本文简介
看完这篇博文，你将了解 pytorch 神经网络的构建、训练、保存、绘图，pytorch xarray加载气象数据构建dataset进行神经网络训练，用来以及一些简单的神经网络训练技巧。
神经网络的构建
神经网络的构建我这里使用的是比较流行的pytorch库，关于此库的下载可以直接百度，选择合适的版本。
根据原文的描述如上图（引用自Ham, YG., Kim, JH. & Luo, JJ. Deep learning for multi-year ENSO forecasts. Nature 573, 568–572 (2019). https://doi.org/10.1038/s41586-019-1559-7），我们使用三层卷积、两层全连接网络。其中代码中的 M_Num , N_Num代表convolutional filters and neurons in the fully connected layer。详细代码如下：
代码语言：javascript复制"""
CNNClass.py
构建神经网络结构的脚本
"""
import torch
import torch.nn as nn
class ConvNetwork(nn.Module):
"""
M_Num , N_Num代表convolutional filters and neurons in the fully connected layer
"""
def __init__(self, M_Num, N_Num):
self.M = M_Num
self.N = N_Num
super().__init__()
self.conv = nn.Sequential(
nn.Conv2d(6, M_Num, kernel_size=(4, 8), padding="same"),
nn.Tanh(),
nn.MaxPool2d(kernel_size=(2, 2), stride=(2, 2)),
nn.Conv2d(M_Num, M_Num, kernel_size=(4, 2), padding="same"),
nn.Tanh(),
nn.MaxPool2d(stride=(2, 2), kernel_size=(2, 2)),
nn.Conv2d(M_Num, M_Num, kernel_size=(4, 2), stride=(1, 1), padding="same"),
nn.Tanh(), )
self.dense = nn.Sequential(
nn.Linear(6 * 18 * M_Num, N_Num),
nn.Linear(N_Num, 23))
def forward(self, InData):
x = self.conv(InData)
x = x.reshape(-1, 6 * 18 * self.M)
x = self.dense(x)
return x
卷积层中的padding使用的方式是 "same",如果你的pytorch不支持此方式，请更新torch版本。
使用nn.Sequential简化神经网络的表达。
forward中x = x.reshape(-1, 6 * 18 * self.M) 语句是用来更改张量形状，链接CNN与FC。
数据加载模块的构建
根据pytorch官方例子Datasets & DataLoaders — PyTorch Tutorials 1.9.0 cu102 documentation,自己写一个dataset要写 init, len, and getitem.三个基础功能，分别对应着数据集初始化（加载文件）、数据集长度、得到对应Index的case。而且数据集可以通过ConcatDataset进行拼接（下面会用到）。所以我们Xarray库和numpy库将我们原来准备的NC文件加载出来。
代码语言：javascript复制"""
DataLoaderFunc.py
数据加载模块,用来加载数据
"""
import torch
import torch.nn
import numpy as np
from torch.utils.data import Dataset, DataLoader
import xarray as xr
CIMPTosLoc = "../TrainData/TosA.nc"
CIMPZosLoc = "../TrainData/ZosA.nc"
CIMPNinoLoc = "../TrainData/Cmip6Nino34I.nc"
OBSTrainSSTALoc = "../TrainData/ersstv5ssta.nc"
OBSTrainSSHALoc = "../TrainData/SODAssha.nc"
OBSTrainNinoLoc = "../TrainData/ersstv5Nino34.nc"
OBSValSSTALoc = "../ValidationData/ersstv5ssta.nc"
OBSValSSHALoc = "../ValidationData/GODASssha.nc"
OBSValNinoLoc = "../ValidationData/ersstv5Nino34.nc"
class ENSODataset(Dataset):
"""
type_ : 需要加载的数据类型，分为
CMIP 加载训练数据
OBSTrain 加载观测数据
OBSVal 加载观测验证数据
T_begin 开始时间
T_end 结束时间
其中，后两个参数只适用于 CMIP数据
"""
def __init__(self, type_, T_begin=None, T_end=None):
super().__init__()
self.Type = type_
if type_ == "CMIP":
self.IniSSTA = xr.open_dataset(CIMPTosLoc)["TosA"].fillna(0) # 去掉 nan
self.IniSSHA = xr.open_dataset(CIMPZosLoc)["ZosA"].fillna(0) # 去掉 nan
self.IniNino = xr.open_dataset(CIMPNinoLoc)["Nino34I"] # 注意开头末尾都为0
if T_begin is not None:
# 选择合适的时间
TimeNeed = (self.IniSSTA["time"].dt.year >= T_begin) & (self.IniSSTA["time"].dt.year <= T_end)
self.IniSSTA = self.IniSSTA[TimeNeed]
self.IniSSHA = self.IniSSHA[TimeNeed]
self.IniNino = self.IniNino[TimeNeed]
self.DataTimeLen = self.IniNino.shape[0]
elif type_ == "OBSTrain":
SSTA = xr.open_dataset(OBSTrainSSTALoc)["ssta"].fillna(0) # 去掉 nan
SSHA = xr.open_dataset(OBSTrainSSHALoc)["ssha"].fillna(0)
TimeNeedSSTA = (SSTA["time"].dt.year >= 1871) & (SSTA["time"].dt.year <= 1973)
TimeNeedSSHA = (SSHA["time"].dt.year >= 1871) & (SSHA["time"].dt.year <= 1973)
self.IniSSTA = SSTA[TimeNeedSSTA]
self.IniSSHA = SSHA[TimeNeedSSHA]
self.IniNino = xr.open_dataset(OBSTrainNinoLoc)["nino34"][TimeNeedSSTA]
self.DataTimeLen = self.IniNino.shape[0]
elif type_ == "OBSVal":
SSTA = xr.open_dataset(OBSValSSTALoc)["ssta"].fillna(0) # 去掉 nan
SSHA = xr.open_dataset(OBSValSSHALoc)["ssha"].fillna(0)
TimeNeedSSTA = (SSTA["time"].dt.year >= 1984) & (SSTA["time"].dt.year <= 2017)
TimeNeedSSHA = (SSHA["time"].dt.year >= 1984) & (SSHA["time"].dt.year <= 2017)
self.IniSSTA = SSTA[TimeNeedSSTA]
self.IniSSHA = SSHA[TimeNeedSSHA]
self.IniNino = xr.open_dataset(OBSTrainNinoLoc)["nino34"][TimeNeedSSTA]
self.DataTimeLen = self.IniNino.shape[0]
else:
raise ValueError("必须是 CMIP,OBSTrain,OBSVal 其中之一")
# 及时检查数据的最大值、最小值
print("数据集名称",type_)
print("数据最大值:", self.IniSSTA.max().item(), self.IniSSHA.max().item(), self.IniNino.max().item())
print("数据最小值:", self.IniSSTA.min().item(), self.IniSSHA.min().item(), self.IniNino.min().item())
def __getitem__(self, index):
"""
index 指的是 一个样本在 数据集中的索引
"""
DataX1 = np.array(self.IniSSTA[index:index 3])
DataX2 = np.array(self.IniSSHA[index:index 3])
DataX = np.concatenate([DataX1, DataX2], axis=0)
DataX = torch.tensor(DataX, dtype=torch.float32)
DataY = np.array(self.IniNino[index 3:index 3 23])
DataY = torch.tensor(DataY, dtype=torch.float32)
return DataX, DataY
def __len__(self):
if self.Type == "CMIP":
return int(self.DataTimeLen - 3 - 23) # 开头三个月、末尾 二十四个月，注意到CMIP最后一个月数据为NAN
else:
return int(self.DataTimeLen - 3 - 23 1)
神经网络训练函数
为了便于日后训练，编写两个函数分别用于验证神经网络和验证神经网络。其中会生成、保存训练日志，训练日志内容有 名称、日期、epoch数、losslist、验证集loss、相关系数、P值list。非常建议大家使用Tensorboard辅助训练。https://www.tensorflow.org/tensorboard
代码语言：javascript复制"""
TrainFuncVal.py
此脚本写用于训练、验证的函数
"""
import torch
import datetime
import scipy.stats as sps
from DataLoaderFunc import ENSODataset
from torch.utils.data import DataLoader
import pickle
from CNNClass import ConvNetwork
def valFunc(Network, DataL=iter(DataLoader(ENSODataset(type_="OBSVal"), batch_size=1000, shuffle=False)),
criterion=torch.nn.MSELoss()):
"""
对神经网络训练的结果进行验证，使用文中给的相关系数和Loss进行验证
Network:需要验证的神经网络
DataL: 需要验证的神经网络数据集加载器，默认为 OBSVal,DataLoader需要一次Load完响相应数据
criterion: 判别器，默认为MSELoss
return
MSELoss , 每个月相关系数 ， 每个相关系数的 P值
"""
inputs, outputs = next(DataL)
# 加载到CPU验证
Network.to("cpu")
Network.eval()
pred = Network.forward(inputs)
LossVal = criterion(pred, outputs).item()
Calpred = pred.T.detach().numpy()
CalOutputs = outputs.T.detach().numpy()
ACCList = []
PList = []
# 对23个月计算相关系数
for index_month in range(23):
acc, p_value = sps.pearsonr(Calpred[index_month], CalOutputs[index_month])
ACCList.append(acc)
PList.append(p_value)
return LossVal, ACCList, PList
def trainFunc(Network, DataL, epochs, optim, saveName, criterion=torch.nn.MSELoss(), device=torch.device("cpu"),
ValMethod=True, gen_log=True, Save=True,
val_data=iter(DataLoader(ENSODataset(type_="OBSVal"), batch_size=1000, shuffle=False))):
"""
此函数对神经网络进行训练、保存
Network:需要验证的神经网络
DataL: 需要训练神经网络数据集加载器
epochs :需要训练的回合数
optim : 神经网络优化器
saveName : 神经网络保存名称
criterion : 判别器 默认torch.nn.MSELoss()
device : 训练的设备 ， 下载了CUDA可以使用GPU训练
ValMethod : 是否需要验证
gen_log : 是否需要产生训练日志
Save ： 神经网络参数是否需要保存
val_data ： 验证数据集
"""
Network.to(device) # 加载网路到设备
Network.train() # 开始训练
lossList = [] # 用于保存训练中的 Loss
for epoch in range(epochs):
for AData, (inputs, outputs) in enumerate(DataL):
# 开始 一个 batch训练
Network.train()
# 加载数据到设备
inputs = inputs.to(device)
outputs = outputs.to(device)
# 优化器 清零
optim.zero_grad()
# 神经网络输出
pred = Network.forward(inputs)
# 计算loss并保存
loss = criterion(pred, outputs)
lossList.append(loss.item())
# 反向传播算法
loss.backward()
optim.step()
# 对一部分训练过程的Loss进行打印
if AData % 5 == 0:
print('epoch={},batch={},损失 = {:.6f}'.format(epoch, AData, loss.item()))
# 获取时间
time_str = datetime.datetime.strftime(datetime.datetime.now(), '%Y-%m-%d %H:%M:%S')
# 保存网络参数
if Save is True:
torch.save(Network.state_dict(), "./NetParam/%s.net" % saveName)
# 神经网络进行检验
if ValMethod is True:
LossVal, ACCList, PList = valFunc(Network, DataL=val_data)
else:
LossVal, ACCList, PList = None, None, None
# 生成并保存日志
if gen_log is True:
# 训练日志包括 名称、日期、epoch数、losslist、验证集loss、相关系数、P值list
SaveDict = {"trainName": saveName, "train_time": time_str, "epoch_num": epochs, "lossList": lossList,
"LossVal": LossVal,
"ACCList": ACCList, "Plist": PList}
saveF = open("TrainRes/%s.pickle" % saveName, "wb")
pickle.dump(SaveDict, saveF)
saveF.close()
图省事，我有写了一个函数来plot神经网络训练过程和验证集技巧。
代码语言：javascript复制"""
FuncPlot.py
函数来plot神经网络训练过程和验证集技巧
"""
import pickle
import numpy as np
import matplotlib.pyplot as plt
def trainPlot(FName):
File = open("TrainRes/" FName ".pickle", "rb")
Dic = pickle.load(File)
fig = plt.figure(figsize=(10, 9))
ax1 = fig.add_subplot(211)
ax1.plot(np.arange(1, 24), Dic["ACCList"], "-o", label="CNN")
ax1.hlines(0.5, 0.5, 23.5)
ax1.set_xlim(0.5, 23.5)
ax1.set_ylim(0, 1)
ax1.set_xlabel("lead time (month)")
ax1.set_ylabel("ACC")
plt.legend()
plt.xticks(np.arange(1, 24, 1))
plt.title(Dic["train_time"], loc="right")
plt.title("ExpName:" Dic["trainName"], loc="left")
ax2 = fig.add_subplot(212)
ax2.plot(Dic["lossList"], label="Train Loss")
ax2.hlines(Dic["LossVal"], 0, len(Dic["lossList"]), label="Val Loss", colors="red")
ax2.set_ylabel("Loss")
ax2.set_xlabel("batch_number")
plt.legend()
plt.savefig("./TrainRes/%s.png" % FName, dpi=300)
plt.show()
注意，上述训练过程如果你有一块好的GPU建议在GPU上训练。而作者本人由于疫情原因困在学校，手头没有合适gpu进行训练，所以在cpu上进行了训练。如果有GPU，需要下载CUDA（一定要对应Pytorch版本)进行训练。则device可以写为
代码语言：javascript复制device = torch.device("cuda:0")
训练过程
训练过程就要轻松很多了，根据我们上述写好的功能，只需要几行，我们就能完成训练。（笔者在这里只用了GFDL一个CMIP6模式数据 观测数据进行了训练，原因是笔者手上没有合适的GPU，不能进行大规模训练，所以训练结果要比原文差）。训练脚本如下
代码语言：javascript复制"""
一个简简单单的训练脚本，训练名叫 Exp1
"""
import torch
from DataLoaderFunc import ENSODataset
from torch.utils.data import DataLoader, ConcatDataset
from CNNClass import ConvNetwork
import TrainFuncVal as TFV
import FuncPlot
import os
# 忽略 matplotlib与torch的一个加载错误。
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
# 初始化网络
Net = ConvNetwork(30, 30)
# 加载数据
DS = ENSODataset("CMIP")
DS1 = ENSODataset("OBSTrain")
# 生成DataLoader
DL = DataLoader(ConcatDataset((DS, DS1)), batch_size=400, shuffle=True)
# 开始训练与验证
TFV.trainFunc(Net, DL, 5, saveName="Exp1", optim=torch.optim.Adam(Net.parameters()))
# 如果需要迁移学习 ， 可以 将Dataset分开加载
# 画出结果
FuncPlot.trainPlot("Exp1")
训练结果如下图，可以看到，只是勉勉强强达到了9个月，与原文水平相距甚远，主要原因是样本不够，再者是我没有用迁移学习方法，而是将CMIP数据与观测数据混着学习(ConcatDataset((DS, DS1)),将观测数据与CMIP数据一起加载)。而且，通过Loss图可以看到，有一点点过拟合(Val loss大于后来的Train loss)。
训练过程可能会报错：
代码语言：javascript复制UserWarning: Named tensors and all their associated APIs are an experimental feature and subject to change. Please do not use them for anything important until they are released as stable. (Triggered internally at ..c10/core/TensorImpl.h:1156.)
return torch.max_pool2d(input, kernel_size, stride, padding, dilation, ceil_mode)
这个报错的原因是pytorch的一个bug，开发者会在下个版本修复。
也可能会报错：
代码语言：javascript复制UserWarning: Using padding='same' with even kernel lengths and odd dilation may require a zero-padded copy of the input be created (Triggered internally at ..atensrcATennativeConvolution.cpp:660.)
return F.conv2d(input, weight, bias, self.stride,
这个的原因是我们使用 padding = "same" ， 由于原网络kernel_size为偶数，不得不在上左多补一行零导致的，对我们的训练没有任何影响。
神经网络的一些训练技巧
使用少量数据预训练模型进行摸底很重要
batch size不是越大越好
Adam能够解决很多疑难杂症，不要一上来就用SGD（SGD有时候真的很好用）
batch_size设置成1有奇效
tensorboard很好用
参数初始化很重要，选择好的初始化方法事半功倍。
结语
本文介绍了神经网络的构建与训练，希望大家能够有所收获。本工作主要完成于2020年秋天、冬天，据此已经过去半年多，有些错误、水平低下的地方在所难免，希望大家多多包涵。
