https://cloud.tencent.com/developer/article/1622038
(强对流天气临近预报)时空序列预测模型—PredRNN(Pytorch)

代码分为3文件：

PredRNN_Cell.py #细胞单元

PredRNN_Model.py #细胞单元堆叠而成的主干模型

PredRNN_Main_Seq2seq_test.py #用于外推的Seq2seq 编码解码

# PredRNN_Cell
import torch.nn as nn
from torch.autograd import Variable
import torch
####################################
#
# 单层，单时间步的PredRNNCell(细胞/单元)，用于构造整个外推模型
# The cell/unit of predrnncell of every layer and time_step, for constructing the entire extrapolation model.
#
####################################
class PredRNNCell(nn.Module):def __init__(self, input_size, input_dim, hidden_dim_m, hidden_dim,kernel_size, bias):super(PredRNNCell, self).__init__()self.height, self.width = input_sizeself.input_dim = input_dimself.hidden_dim = hidden_dimself.hidden_dim_m = hidden_dim_m   #  hidden of Mself.kernel_size = kernel_sizeself.padding = kernel_size[0] // 2, kernel_size[1] // 2self.bias = bias###################################################################################### 相应符号可对应参照论文# Corresponding symbols can correspond to reference paper# conv_h_c for gt, it, ft# conv_m for gt', it', ft'# conv_o for ot# self.conv_h_next for Htself.conv_h_c = nn.Conv2d(in_channels=self.input_dim + self.hidden_dim,out_channels=3 * self.hidden_dim,kernel_size=self.kernel_size,padding=self.padding,bias=self.bias)self.conv_m = nn.Conv2d(in_channels=self.input_dim + self.hidden_dim_m,out_channels=3 * self.hidden_dim_m,kernel_size=self.kernel_size,padding=self.padding,bias=self.bias)self.conv_o = nn.Conv2d(in_channels=self.input_dim + self.hidden_dim * 2 + self.hidden_dim_m,out_channels=self.hidden_dim,kernel_size=self.kernel_size,padding=self.padding,bias=self.bias)self.conv_h_next = nn.Conv2d(in_channels=self.hidden_dim + self.hidden_dim_m,out_channels=self.hidden_dim,kernel_size=1,bias=self.bias)def forward(self, input_tensor, cur_state, cur_state_m):h_cur, c_cur= cur_state  #cur = Current input of H and Ch_cur_m = cur_state_m #cur = Current input of mcombined_h_c = torch.cat([input_tensor,h_cur], dim=1)combined_h_c = self.conv_h_c(combined_h_c)cc_i, cc_f, cc_g = torch.split(combined_h_c, self.hidden_dim, dim=1)combined_m = torch.cat([input_tensor,  h_cur_m], dim=1)combined_m = self.conv_m(combined_m)cc_i_m, cc_f_m, cc_g_m = torch.split(combined_m, self.hidden_dim_m, dim=1)i = torch.sigmoid(cc_i)f = torch.sigmoid(cc_f)g = torch.tanh(cc_g)c_next = f * c_cur + i * gi_m = torch.sigmoid(cc_i_m)f_m = torch.sigmoid(cc_f_m)g_m = torch.tanh(cc_g_m)h_next_m = f_m * h_cur_m + i_m * g_mcombined_o = torch.cat([input_tensor, h_cur, c_next, h_next_m], dim=1)combined_o = self.conv_o(combined_o)o = torch.sigmoid(combined_o)h_next = torch.cat([c_next, h_next_m], dim=1)h_next = self.conv_h_next(h_next)h_next = o * torch.tanh(h_next)return h_next, c_next, h_next_m####################################### 用于在t=0时刻时初始化H,C,M# For initializing H,C,M at t=0######################################def init_hidden(self, batch_size):return (Variable(torch.zeros(batch_size, self.hidden_dim, self.height, self.width)).cuda(),Variable(torch.zeros(batch_size, self.hidden_dim, self.height, self.width)).cuda())

#PredRNN_Model
import torch.nn as nn
import torch
import numpy as np
from torch.autograd import Variable
from PredRNN_Cell import PredRNNCell
##############################################
#
# 构造PredRNN
# Construct PredRNN
#
##############################################
class PredRNN(nn.Module):def __init__(self, input_size, input_dim, hidden_dim, hidden_dim_m, kernel_size, num_layers,batch_first=False, bias=True):super(PredRNN, self).__init__()self._check_kernel_size_consistency(kernel_size)kernel_size = self._extend_for_multilayer(kernel_size, num_layers) # 按照层数来扩充 卷积核尺度/可自定义hidden_dim = self._extend_for_multilayer(hidden_dim, num_layers) # 按照层数来扩充 LSTM单元隐藏层维度/可自定义hidden_dim_m = self._extend_for_multilayer(hidden_dim_m, num_layers)  # M的单元应保持每层输入和输出的一致性.if not len(kernel_size) == len(hidden_dim) == num_layers:  # 判断相应参数的长度是否与层数相同raise ValueError('Inconsistent list length.')self.height, self.width = input_sizeself.input_dim = input_dimself.hidden_dim = hidden_dimself.hidden_dim_m = hidden_dim_mself.kernel_size = kernel_sizeself.num_layers = num_layersself.batch_first = batch_firstself.bias = biascell_list = []for i in range(0, self.num_layers):if i == 0:    # 0 时刻， 图片的输入即目前实际输入cur_input_dim = self.input_dimelse:cur_input_dim = self.hidden_dim[i - 1]  # 非0时刻，输入的维度为上一层的输出#Cell_list.appenda为堆叠层操作cell_list.append(PredRNNCell(input_size=(self.height, self.width),input_dim=cur_input_dim,hidden_dim=self.hidden_dim[i],hidden_dim_m=self.hidden_dim_m[i],kernel_size=self.kernel_size[i],bias=self.bias).cuda())self.cell_list = nn.ModuleList(cell_list)#Cell_list进行Model化def forward(self, input_tensor, hidden_state=False, hidden_state_m=False):if self.batch_first is False:input_tensor = input_tensor.permute(1, 0, 2, 3, 4)if hidden_state is not False:hidden_state = hidden_stateelse:   #如果没有输入自定义的权重，就以0元素来初始化hidden_state = self._init_hidden(batch_size=input_tensor.size(0))if hidden_state_m is False:h_m = Variable(torch.zeros(input_tensor.shape[0], self.hidden_dim_m[0],input_tensor.shape[3], input_tensor.shape[4]),requires_grad=True).cuda()else:h_m = hidden_state_mlayer_output_list = []  #记录输出layer_output_list_m = []  # 记录每一层mlayer_output_list_c = []  # 记录每一层clast_state_list = []  #记录最后一个状态layer_output_list_m = []  # 记录最后一个mlast_state_list_m = []  # 记录最后一个mseq_len = input_tensor.size(1)   #第二个时间序列，3cur_layer_input_1 = input_tensor #x方向上的输入all_layer_out = []for t in range(seq_len):concat=[]output_inner_c = []  # 记录输出的coutput_inner = []  #记录输出的coutput_inner_m = []  # 记录输出的moutput_inner_h_c=[] # 记录输出的h 和ch0 = cur_layer_input_1[:, t, :, :, :]  #确定layer = 1 时的输入,如雷达回波图等矩阵信息for layer_idx in range(self.num_layers): # 由于M在layer上传递,所以优先考虑layer上的传递if t == 0:  # 由于M在layer上传递，所以要区分t=0(此时m初始化)h, c = hidden_state[layer_idx]  # h和c来自于初始化/自定义h, c, h_m = self.cell_list[layer_idx](input_tensor=h0,cur_state=[h, c], cur_state_m=h_m)#经过一个cell/units输出的h,c,moutput_inner_c.append(c) #记录输出的c进行output_inner.append(h)output_inner_m.append(h_m)output_inner_h_c.append([h,c])h0=helse:h = cur_layer_input[layer_idx]c = cur_layer_input_c[layer_idx]h, c, h_m = self.cell_list[layer_idx](input_tensor=h0,cur_state=[h, c], cur_state_m=h_m)output_inner_c.append(c)output_inner.append(h)output_inner_m.append(h_m)output_inner_h_c.append([h, c])h0 = hcur_layer_input = output_inner#记录某个t，全部layer的输出hcur_layer_input_c = output_inner_c#记录某个t，全部layer的输出ccur_layer_input_m = output_inner_m#记录某个t，全部layer的输出malllayer_output = torch.cat(output_inner, dim=1) #把某个t时刻每个隐藏层的输出进行堆叠,以便于在解码层参照Convlstm使用1x1卷积得到输出all_layer_out.append(alllayer_output)#记录每个t时刻,所有隐藏层输出的h,以便于在解码层参照Convlstm使用1x1卷积得到输出per_time_all_layer_stack_out=torch.stack(all_layer_out, dim=1)#记录每个t时刻,所有隐藏层输出的h,以便于在解码层参照Convlstm使用1x1卷积得到输出layer_output_list.append(h)# 记录每一个t得到的最后layer的输出hlast_state_list.append([h, c])#记录每一个t得到的最后layer的输出h,Clast_state_list_m.append(h_m)#记录每一个t得到的最后layer的输出m#按层对最后一层的H和C进行扩展# ↓↓↓↓↓↓↓↓↓全部t时刻最后layer的输出h# ↓↓↓↓↓↓↓↓↓最后t时刻全部layer的输出h和c# ↓↓↓↓↓↓↓↓↓全部t时刻最后layer的输出m/t+1时刻0 layer的输入m# ↓↓↓↓↓↓↓↓↓全部时刻全部layer的h在隐藏层维度上的总和，hidden_dim = [7,1],则输出channels = 8return torch.stack(layer_output_list, dim=1),\output_inner_h_c,\torch.stack(last_state_list_m, dim=0),\per_time_all_layer_stack_outdef _init_hidden(self, batch_size):init_states = []for i in range(self.num_layers):init_states.append(self.cell_list[i].init_hidden(batch_size))return init_states@staticmethoddef _check_kernel_size_consistency(kernel_size):if not (isinstance(kernel_size, tuple) or(isinstance(kernel_size, list) and all([isinstance(elem, tuple) for elem in kernel_size]))):raise ValueError('`kernel_size` must be tuple or list of tuples')@staticmethoddef _extend_for_multilayer(param, num_layers):if not isinstance(param, list):param = [param] * num_layersreturn param

#PredRNN_Seq2Seq
from PredRNN_Model import PredRNN
import torch.optim as optim
import torch.nn as nn
import matplotlib.pyplot as plt
import numpy as np
from torch.autograd import Variable
import torch
from PIL import Image
from torch.utils.data import Dataset, DataLoader
import os
import timeinput=torch.rand(1,1,1,100,100).cuda()   # Batch_size , time_step, channels, hight/width, width/hight
target=torch.rand(1,1,1,100,100).cuda()   # Batch_size , time_step, channels, hight/width, width/hightclass PredRNN_enc(nn.Module):def __init__(self):super(PredRNN_enc, self).__init__()self.pred1_enc=PredRNN(input_size=(100,100),input_dim=1,hidden_dim=[7, 1],hidden_dim_m=[7, 7],kernel_size=(7, 7),num_layers=2,batch_first=True,bias=True).cuda()def forward(self,enc_input):_, layer_h_c, all_time_h_m, _ = self.pred1_enc(enc_input)return layer_h_c, all_time_h_m
class PredRNN_dec(nn.Module):def __init__(self):super(PredRNN_dec, self).__init__()self.pred1_dec=PredRNN(input_size=(100,100),input_dim=1,hidden_dim=[7, 1],hidden_dim_m=[7, 7],kernel_size=(7, 7),num_layers=2,batch_first=True,bias=True).cuda()self.relu = nn.ReLU()def forward(self,dec_input,enc_hidden,enc_h_m):out, layer_h_c, last_h_m, _ = self.pred1_dec(dec_input,enc_hidden,enc_h_m)out = self.relu(out)return out, layer_h_c, last_h_m
enc=PredRNN_enc().cuda()
dec=PredRNN_dec().cuda()import itertools
loss_fn=nn.MSELoss()
position=0
optimizer=optim.Adam(itertools.chain(enc.parameters(), dec.parameters()),lr=0.001)
for epoch in range(1000):loss_total=0enc_hidden, enc_h_m = enc(input)for i in range(input.shape[1]):optimizer.zero_grad()out, layer_h_c, last_h_m = dec(input[:,i:i+1,:,:,:], enc_hidden, enc_h_m[-1])loss=loss_fn(out, target[:,i:i+1,:,:,:])loss_total+=lossenc_hidden = layer_h_cenc_h_m = last_h_mloss_total=loss_total/input.shape[1]loss_total.backward()optimizer.step()print(epoch,epoch,loss_total)