|
- from functools import partial
- from collections import OrderedDict
- import torch
- import torch.nn as nn
- import torch.nn.functional as F
- import numpy as np
- from timm.models.layers import DropPath, to_2tuple, trunc_normal_, to_3tuple
- import torch.fft
- from params import get_args
- from torch.utils.checkpoint import checkpoint_sequential
- import random
- from torch.utils.data import Dataset, DataLoader
- from tqdm import tqdm
- from sklearn.metrics import mean_squared_error
- from sklearn.preprocessing import MinMaxScaler
-
- class Mlp(nn.Module):
- def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
- super().__init__()
- out_features = out_features or in_features
- hidden_features = hidden_features or in_features
- self.fc1 = nn.Linear(in_features, hidden_features)
- self.act = act_layer()
- # self.fc2 = nn.Linear(hidden_features, out_features)
- self.fc2 = nn.AdaptiveAvgPool1d(out_features)
- self.drop = nn.Dropout(drop)
-
- def forward(self, x):
- x = self.fc1(x)
- x = self.act(x)
- x = self.drop(x)
- x = self.fc2(x)
- x = self.drop(x)
- return x
-
-
- class AdaptiveFourierNeuralOperator(nn.Module):
- def __init__(self, dim, h=14, w=8):
- super().__init__()
- args = get_args()
- self.hidden_size = dim
- self.h = h
- self.w = w
-
- self.num_blocks = args.fno_blocks
- self.block_size = self.hidden_size // self.num_blocks
- assert self.hidden_size % self.num_blocks == 0
-
- self.scale = 0.02
- self.w1 = torch.nn.Parameter(self.scale * torch.randn(2, self.num_blocks, self.block_size, self.block_size))
- self.b1 = torch.nn.Parameter(self.scale * torch.randn(2, self.num_blocks, self.block_size))
- self.w2 = torch.nn.Parameter(self.scale * torch.randn(2, self.num_blocks, self.block_size, self.block_size))
- self.b2 = torch.nn.Parameter(self.scale * torch.randn(2, self.num_blocks, self.block_size))
- self.relu = nn.ReLU()
-
- if args.fno_bias:
- self.bias = nn.Conv1d(self.hidden_size, self.hidden_size, 1)
- else:
- self.bias = None
-
- self.softshrink = args.fno_softshrink
-
- def multiply(self, input, weights):
- return torch.einsum('...bd,bdk->...bk', input, weights)
-
- def forward(self, x):
- B, N, C = x.shape
-
- if self.bias:
- bias = self.bias(x.permute(0, 2, 1)).permute(0, 2, 1)
- else:
- bias = torch.zeros(x.shape, device=x.device)
-
- x = x.reshape(B, self.h, self.w, C)
- x = torch.fft.rfft2(x, dim=(1, 2), norm='ortho')
- x = x.reshape(B, x.shape[1], x.shape[2], self.num_blocks, self.block_size)
-
- x_real = F.relu(self.multiply(x.real, self.w1[0]) - self.multiply(x.imag, self.w1[1]) + self.b1[0],
- inplace=True)
- x_imag = F.relu(self.multiply(x.real, self.w1[1]) + self.multiply(x.imag, self.w1[0]) + self.b1[1],
- inplace=True)
- x_real = self.multiply(x_real, self.w2[0]) - self.multiply(x_imag, self.w2[1]) + self.b2[0]
- x_imag = self.multiply(x_real, self.w2[1]) + self.multiply(x_imag, self.w2[0]) + self.b2[1]
-
- x = torch.stack([x_real, x_imag], dim=-1)
- x = F.softshrink(x, lambd=self.softshrink) if self.softshrink else x
-
- x = torch.view_as_complex(x)
- x = x.reshape(B, x.shape[1], x.shape[2], self.hidden_size)
- x = torch.fft.irfft2(x, s=(self.h, self.w), dim=(1, 2), norm='ortho')
- x = x.reshape(B, N, C)
-
- return x + bias
-
-
- class Block(nn.Module):
- def __init__(self, dim, mlp_ratio=4., drop=0., drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, h=14, w=8):
- super().__init__()
- args = get_args()
- self.norm1 = norm_layer(dim)
- self.filter = AdaptiveFourierNeuralOperator(dim, h=h, w=w)
-
- self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
- self.norm2 = norm_layer(dim)
- mlp_hidden_dim = int(dim * mlp_ratio)
- self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
-
- self.double_skip = args.double_skip
-
- def forward(self, x):
- residual = x
- x = self.norm1(x)
- x = self.filter(x)
-
- if self.double_skip:
- x += residual
- residual = x
-
- x = self.norm2(x)
- x = self.mlp(x)
- x = self.drop_path(x)
- x += residual
- return x
-
-
- class PatchEmbed(nn.Module):
- def __init__(self, img_size=None, patch_size=8, in_chans=20, embed_dim=768):
- super().__init__()
-
- if img_size is None:
- raise KeyError('img is None')
-
- patch_size = to_3tuple(patch_size)
-
- num_patches =1 * (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0])
- self.img_size = img_size
- self.patch_size = patch_size
- self.num_patches = num_patches
- self.proj = nn.Conv3d(in_chans, embed_dim, kernel_size=(8,8,8), stride=(8,8,8))
-
- def forward(self, x):
- # print('x.shape:{}'.format(x.shape)) #x.shape:torch.Size([10, 7, 20, 40, 208])
- B, T, C, H, W = x.shape
- x = x.permute(0,2,1,3,4)
- # print('x.shape:{}'.format(x.shape))
- # FIXME look at relaxing size constraints
- assert H == self.img_size[0] and W == self.img_size[1] and T == self.img_size[2], f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
- x = self.proj(x) # x.shape:torch.Size([10, 768, 8, 5, 26])
- # print('x1.shape:{}'.format(x.shape)) #x.shape:torch.Size([10, 130, 768])
- x = x.flatten(3)
- # print('x.shape:{}'.format(x.shape)) # x.shape:torch.Size([10, 768, 8, 130])
- x = x.flatten(2)
- # print('x.shape:{}'.format(x.shape)) # x.shape:torch.Size([10, 768, 1040])
- x = x.transpose(1,2)
- # print('x.shape:{}'.format(x.shape)) # x.shape:torch.Size([10, 1040, 768])
- return x
-
-
- class AFNONet(nn.Module):
- def __init__(self, img_size=None, patch_size=8, in_chans=2, out_chans=2, embed_dim=768, depth=12, mlp_ratio=4.,
- uniform_drop=False, drop_rate=0., drop_path_rate=0., norm_layer=None, dropcls=0):
- super().__init__()
-
- if img_size is None:
- img_size = [40, 200, 10]
-
- self.embed_dim = embed_dim
- norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6)
-
- self.patch_embed = PatchEmbed(img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)
- num_patches = self.patch_embed.num_patches
-
- # print('num_patch:{}'.format(num_patches)) # 130 #num_patch:1750
- self.pos_embed = nn.Parameter(torch.zeros(1, num_patches, embed_dim)) # 可学习的参数 “pos_embed# ”
-
- self.pos_drop = nn.Dropout(p=drop_rate)
-
- self.t = 1
- self.h = img_size[0] // patch_size
- self.w = (img_size[1] // patch_size) * self.t
-
-
- if uniform_drop:
- dpr = [drop_path_rate for _ in range(depth)] # stochastic depth decay rule
- else:
- dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)] # stochastic depth decay rule
-
- self.blocks = nn.ModuleList([Block(dim=embed_dim, mlp_ratio=mlp_ratio, drop=drop_rate, drop_path=dpr[i],
- norm_layer=norm_layer, h=self.h, w=self.w) for i in range(depth)])
- self.norm = norm_layer(embed_dim)
-
- # Representation layer
- # self.num_features = out_chans * img_size[0] * img_size[1]
- # self.representation_size = self.num_features * 8
- # self.pre_logits = nn.Sequential(OrderedDict([
- # ('fc', nn.Linear(embed_dim, self.representation_size)),
- # ('act', nn.Tanh())
- # ]))
- self.pre_logits = nn.Sequential(OrderedDict([
- ('conv1', nn.ConvTranspose3d(embed_dim, out_chans * 16, kernel_size=(2, 2, 1), stride=(2, 2, 1))),
- # 将输入特征图的高度和宽度分别扩大两倍,深度保持不变 (10, 768, 5, 26, 8) --> (10, 320, 10, 52, 8)
- ('act1', nn.Tanh()), # 形状不变 (10, 320, 10, 52, 8) --> (10, 320, 10, 52, 8)
- ('conv2', nn.ConvTranspose3d(out_chans * 16, out_chans * 4, kernel_size=(2, 2, 1), stride=(2, 2, 1))),
- # 将输入特征图的高度、宽度和深度分别扩大两倍 (10, 320, 10, 52, 8) --> (10, 80, 20, 104, 8)
- ('act2', nn.Tanh()) # 形状不变 (10, 80, 20, 104, 8) --> (10, 80, 20, 104, 8)
- ]))
-
- # Generator head
- # self.head = nn.Linear(self.representation_size, self.num_features)
- self.head = nn.ConvTranspose3d(out_chans * 4, out_chans, kernel_size=(2, 2, 1), stride=(2, 2, 1))
-
- if dropcls > 0:
- print('dropout %.2f before classifier' % dropcls)
- self.final_dropout = nn.Dropout(p=dropcls)
- else:
- self.final_dropout = nn.Identity()
-
- trunc_normal_(self.pos_embed, std=.02)
- self.apply(self._init_weights)
-
- def _init_weights(self, m):
- if isinstance(m, nn.Linear):
- trunc_normal_(m.weight, std=.02)
- if isinstance(m, nn.Linear) and m.bias is not None:
- nn.init.constant_(m.bias, 0)
- elif isinstance(m, nn.LayerNorm):
- nn.init.constant_(m.bias, 0)
- nn.init.constant_(m.weight, 1.0)
-
- @torch.jit.ignore
- def no_weight_decay(self):
- return {'pos_embed', 'cls_token'}
-
- def forward_features(self, x):
- # print('x.shape:{}'.format(x.shape)) # x.shape:torch.Size([10, 20, 41, 210]) 输入数据
- B = x.shape[0]
- x = self.patch_embed(x)
- # print('x.shape:{}'.format(x.shape)) # ([10, 8, 20, 40, 208]) ---> ([10, 1040, 768]) # 输入为[B, C, H, W] 输出为 [b,num_patches, embed_dim]
- # num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0])
- # patch_size=8 [41,210]--->{210//8} * (41 //8) = 26 * 5 = 130
- x += self.pos_embed # x.shape:torch.Size([10, 130, 768]) 可学习的位置编码
- # print('x.shape:{}'.format(x.shape))
- x = self.pos_drop(x) # x.shape:torch.Size([10, 130, 768]) 随即正则化失活 dropout
- # print('x1.shape:{}'.format(x.shape)) #ok
-
- if not get_args().checkpoint_activations:
- for blk in self.blocks:
- x = blk(x)
- else:
- x = checkpoint_sequential(self.blocks, 4, x)
-
- # print('x.shape:{}'.format(x.shape)) # x.shape:torch.Size([10, 1040, 768])
- x = self.norm(x).transpose(1, 2) # x.shape:torch.Size([10, 768, 1040])
- x = x.reshape(-1,self.embed_dim, self.h, int(self.w/self.t), self.t) # x.shape:torch.Size([10, 768, 5, 26, 8])
- return x # img_size = [41,210] h = 41//8 = 5 w = 210//8 = 26
-
- def forward(self, x):
- # print('x.shape:{}'.format(x.shape)) # 输入x为 ([10, 20, 41, 210])
- x = self.forward_features(x) # ([10, 7, 20, 41, 210]) ---》 ([10, 768, 5, 26, 8])
- # print('x.shape:{}'.format(x.shape)) # x.shape:torch.Size([10, 768, 5, 25, 1])
-
- # x = x.permute(0,2,3,4,1) #torch.Size([10, 5, 25, 1, 768])
- # B, H, W, T, embid, = x.shape
- # x = x.reshape(B,H*W*T,embid)
- # # print(x.shape) #torch.Size([10, 125, 768])
- # x = self.norm(x)
- # # print(x.shape) #torch.Size([10, 125, 768])
- # x = x.permute(0,2,1)
- # B, embid, C = x.shape
- # x = x.reshape(B,embid,5,25,1)
-
-
- x = self.final_dropout(x) # ([10, 768, 5, 26, 8]) ---》 ([10, 768, 5, 26, 8])
-
- x = self.pre_logits(x) # ([10, 768, 5, 26, 8])---》 x.shape:torch.Size([10, 80, 20, 104, 8])
-
- x = self.head(x) # ([10, 80, 20, 104]) ---》 ([10, 20, 40, 208]) 都变为原来的2倍
- # print('x.shape:{}'.format(x.shape))
-
- x = x.permute(0,1,4,2,3)
- x = torch.squeeze(x)
- # print(x.shape)
-
- return x
-
-
- # if __name__ == '__main__':
- # a = torch.randn(10, 10, 12, 40, 200)
- # net = AFNONet()
- # b = net(a)
- # print(b.shape) # torch.Size([10, 20, 40, 208])
-
-
-
- def setup_seed(seed):
- torch.manual_seed(seed)
- torch.cuda.manual_seed_all(seed)
- np.random.seed(seed)
- random.seed(seed)
- torch.backends.cudnn.deterministic = True
-
-
- # 设置随机数种子
- setup_seed(1)
-
- #需要 mld u v sss temp 降水 蒸发 混合层下的盐度
-
- data = np.load(r'/dataset/10day_for_14day_all_variables_surface_pacific_10_19_Qfenjie.npz')
- print(data.files)
- # 2000-01-01 ---> 2019-12-31
- surface_latent_heat_flux = data['surface_latent_heat_flux'][:] #(7305, 10, 40, 200)
- surface_sensible_heat_flux = data['surface_sensible_heat_flux'][:] # (9851, 10, 40, 200)
- surface_net_radiation = data['surface_net_radiation'][:] # (9851, 10, 40, 200)
- evaporation = data['evaporation'][:]
- total_precipitation = data['total_precipitation'][:]
- mld = data['mld'][:]
- sst_surface = data['sst_surface'][:]
- sss_surface = data['sss_surface'][:]
- u_surface = data['u_surface'][:]
- v_surface = data['v_surface'][:]
- T_d = data['T_d'][:]
- S_d = data['S_d'][:]
- u_d = data['u_d'][:]
- v_d = data['v_d'][:]
- # xx = data['xx'][:,:,:40,:200]
- # yy = data['yy'][:,:,:40,:200]
- sst_surface_label = data['sst_surface_label'][:,0,:,:]
- sss_surface_label = data['sss_surface_label'][:,0,:,:]
-
-
- #归一化
-
- # scaler = MinMaxScaler()
- # surface_latent_heat_flux1 = surface_latent_heat_flux.reshape(-1,1)
- # surface_sensible_heat_flux1 = surface_sensible_heat_flux.reshape(-1,1)
- # surface_net_radiation1 = surface_net_radiation.reshape(-1,1)
- # evaporation1 = evaporation.reshape(-1,1)
- # total_precipitation1 = total_precipitation.reshape(-1,1)
- # mld1 = mld.reshape(-1,1)
- # sst_surface1 = sst_surface.reshape(-1,1)
- # sss_surface1 = sss_surface.reshape(-1,1)
- # u_surface1 = u_surface.reshape(-1,1)
- # v_surface1 = v_surface.reshape(-1,1)
- # T_d1 = T_d.reshape(-1,1)
- # S_d1 = S_d.reshape(-1,1)
- # u_d1 = u_d.reshape(-1,1)
- # v_d1 = v_d.reshape(-1,1)
- # sst_surface_label1 = sst_surface_label.reshape(-1,1)
- # sss_surface_label1 = sss_surface_label.reshape(-1,1)
-
-
- # surface_latent_heat_flux = scaler.fit_transform(surface_latent_heat_flux1).reshape(surface_latent_heat_flux.shape)
- # surface_sensible_heat_flux = scaler.fit_transform(surface_sensible_heat_flux1).reshape(surface_sensible_heat_flux.shape)
- # surface_net_radiation1 = scaler.fit_transform(surface_net_radiation1).reshape(surface_net_radiation.shape)
- # evaporation = scaler.fit_transform(evaporation1).reshape(evaporation.shape)
- # total_precipitation = scaler.fit_transform(total_precipitation1).reshape(total_precipitation.shape)
- # mld = scaler.fit_transform(mld1).reshape(mld.shape)
- # sst_surface = scaler.fit_transform(sst_surface1).reshape(sst_surface.shape)
- # sss_surface = scaler.fit_transform(sss_surface1).reshape(sss_surface.shape)
- # u_surface = scaler.fit_transform(u_surface1).reshape(u_surface.shape)
- # v_surface = scaler.fit_transform(v_surface1).reshape(v_surface.shape)
- # T_d = scaler.fit_transform(T_d1).reshape(T_d.shape)
- # S_d = scaler.fit_transform(S_d1).reshape(S_d.shape)
- # u_d = scaler.fit_transform(u_d1).reshape(u_d.shape)
- # v_d = scaler.fit_transform(v_d1).reshape(v_d.shape)
- # sst_surface_label = scaler.fit_transform(sst_surface_label1).reshape(sst_surface_label.shape)
- # sss_surface_label = scaler.fit_transform(sss_surface_label1).reshape(sss_surface_label.shape)
-
- print(sst_surface)
-
- # print(sss_surface_label.shape)
-
- #(7305,10,40,200)
-
- train_size = 1312
- valid_size = 1760 #前20% 作为验证 剩下的20%的作为测试
- # surface_latent_heat_flux = surface_latent_heat_flux + surface_sensible_heat_flux + surface_net_radiation
-
- surface_latent_heat_flux = surface_latent_heat_flux.reshape(-1, 10, 1, 40, 200)
- surface_latent_heat_flux = torch.Tensor(surface_latent_heat_flux)
- surface_latent_heat_flux_train = surface_latent_heat_flux[0:train_size,:,:,:,:]
- surface_latent_heat_flux_valid = surface_latent_heat_flux[train_size:valid_size,:,:,:,:]
- # surface_latent_heat_flux_test = surface_latent_heat_flux[valid_size:,:,:,:,:]
-
-
- surface_sensible_heat_flux = surface_sensible_heat_flux.reshape(-1, 10, 1, 40, 200)
- surface_sensible_heat_flux = torch.Tensor(surface_sensible_heat_flux)
- surface_sensible_heat_flux_train = surface_sensible_heat_flux[0:train_size,:,:,:,:]
- surface_sensible_heat_flux_valid = surface_sensible_heat_flux[train_size:valid_size,:,:,:,:]
- # # # surface_sensible_heat_flux_test = surface_sensible_heat_flux[valid_size:,:,:,:,:]
-
-
- surface_net_radiation = surface_net_radiation.reshape(-1, 10, 1, 40, 200)
- surface_net_radiation = torch.Tensor(surface_net_radiation)
- surface_net_radiation_train = surface_net_radiation[0:train_size,:,:,:,:]
- surface_net_radiation_valid = surface_net_radiation[train_size:valid_size,:,:,:,:]
- # # surface_net_radiation_test = surface_net_radiation[valid_size:,:,:,:,:]
-
-
- evaporation = evaporation.reshape(-1, 10, 1, 40, 200)
- evaporation = torch.Tensor(evaporation)
- evaporation_train = evaporation[0:train_size,:,:,:,:]
- evaporation_valid = evaporation[train_size:valid_size,:,:,:,:]
- # evaporation_test = evaporation[valid_size:,:,:,:,:]
-
-
- total_precipitation = total_precipitation.reshape(-1, 10, 1, 40, 200)
- total_precipitation = torch.Tensor(total_precipitation)
- total_precipitation_train = total_precipitation[0:train_size,:,:,:,:]
- total_precipitation_valid = total_precipitation[train_size:valid_size,:,:,:,:]
- # total_precipitation_test = total_precipitation[valid_size:,:,:,:,:]
-
-
- mld = mld.reshape(-1, 10, 1, 40, 200)
- mld = torch.Tensor(mld)
- mld_train = mld[0:train_size,:,:,:,:]
- mld_valid = mld[train_size:valid_size,:,:,:,:]
- # mld_test = mld[valid_size:,:,:,:,:]
-
-
- sst_surface = sst_surface.reshape(-1, 10, 1, 40, 200)
- sst_surface = torch.Tensor(sst_surface)
- sst_surface_train = sst_surface[0:train_size,:,:,:,:]
- sst_surface_valid = sst_surface[train_size:valid_size,:,:,:,:]
- # sst_surface_test = sst_surface[valid_size:,:,:,:,:]
-
-
- sss_surface = sss_surface.reshape(-1, 10, 1, 40, 200)
- sss_surface = torch.Tensor(sss_surface)
- sss_surface_train = sss_surface[0:train_size,:,:,:,:]
- sss_surface_valid = sss_surface[train_size:valid_size,:,:,:,:]
- # sss_surface_test = sss_surface[valid_size:,:,:,:,:]
-
-
- u_surface = u_surface.reshape(-1, 10, 1, 40, 200)
- u_surface = torch.Tensor(u_surface)
- u_surface_train = u_surface[0:train_size,:,:,:,:]
- u_surface_valid = u_surface[train_size:valid_size,:,:,:,:]
- # u_surface_test = u_surface[valid_size:,:,:,:,:]
-
-
- v_surface = v_surface.reshape(-1, 10, 1, 40, 200)
- v_surface = torch.Tensor(v_surface)
- v_surface_train = v_surface[0:train_size,:,:,:,:]
- v_surface_valid = v_surface[train_size:valid_size,:,:,:,:]
- # # v_surface_test = v_surface[valid_size:,:,:,:,:]
-
-
- T_d = T_d.reshape(-1, 10, 1, 40, 200)
- T_d = torch.Tensor(T_d)
- T_d_train = T_d[0:train_size,:,:,:,:]
- T_d_valid = T_d[train_size:valid_size,:,:,:,:]
- # T_d_test = T_d[valid_size:,:,:,:,:]
-
-
- S_d = S_d.reshape(-1, 10, 1, 40, 200)
- S_d = torch.Tensor(S_d)
- S_d_train = S_d[0:train_size,:,:,:,:]
- S_d_valid = S_d[train_size:valid_size,:,:,:,:]
- # S_d_test = S_d[valid_size:,:,:,:,:]
-
-
- u_d = u_d.reshape(-1, 10, 1, 40, 200)
- u_d = torch.Tensor(u_d)
- u_d_train = u_d[0:train_size,:,:,:,:]
- u_d_valid = u_d[train_size:valid_size,:,:,:,:]
- # u_d_test = u_d[valid_size:,:,:,:,:]
-
-
- v_d = v_d.reshape(-1, 10, 1, 40, 200)
- v_d = torch.Tensor(v_d)
- v_d_train = v_d[0:train_size,:,:,:,:]
- v_d_valid = v_d[train_size:valid_size,:,:,:,:]
- # # v_d_test = v_d[valid_size:,:,:,:,:]
-
- # xx = xx.reshape(-1, 10, 1, 40, 200)
- # xx = torch.Tensor(xx)
- # xx_train = xx[0:train_size,:,:,:,:]
- # xx_valid = xx[train_size:valid_size,:,:,:,:]
- # # xx_test = xx[valid_size:,:,:,:,:]
-
-
- # yy = yy.reshape(-1, 10, 1, 40, 200)
- # yy = torch.Tensor(yy)
- # yy_train = yy[0:train_size,:,:,:,:]
- # yy_valid = yy[train_size:valid_size,:,:,:,:]
- # yy_test = yy[valid_size:,:,:,:,:]
-
-
-
- # train_data = torch.cat((surface_latent_heat_flux_train, surface_sensible_heat_flux_train, surface_net_radiation_train,
- # evaporation_train, total_precipitation_train, mld_train, sst_surface_train, sss_surface_train,
- # u_surface_train, v_surface_train, T_d_train, S_d_train, u_d_train, v_d_train,
- # xx_train, yy_train), dim=2) # train_data.shape:torch.Size([5920, 10, 16, 40, 200])
-
- # valid_data = torch.cat((surface_latent_heat_flux_valid, surface_sensible_heat_flux_valid, surface_net_radiation_valid,
- # evaporation_valid, total_precipitation_valid, mld_valid, sst_surface_valid, sss_surface_valid,
- # u_surface_valid, v_surface_valid, T_d_valid, S_d_valid, u_d_valid, v_d_valid,
- # xx_valid, yy_valid), dim=2)
-
- # test_data = torch.cat((surface_latent_heat_flux_test, surface_sensible_heat_flux_test, surface_net_radiation_test,
- # evaporation_test, total_precipitation_test, mld_test, sst_surface_test, sss_surface_test,
- # u_surface_test, v_surface_test, T_d_test, S_d_test, u_d_test, v_d_test,
- # xx_test, yy_test), dim=2)
-
-
- # train_data = torch.cat((surface_latent_heat_flux_train, surface_sensible_heat_flux_train, surface_net_radiation_train,
- # evaporation_train, total_precipitation_train, mld_train, sst_surface_train, sss_surface_train,
- # u_surface_train, v_surface_train, T_d_train, S_d_train, u_d_train,v_d_train
- # ), dim=2) # train_data.shape:torch.Size([5920, 10, 16, 40, 200])
-
- # valid_data = torch.cat((surface_latent_heat_flux_valid, surface_sensible_heat_flux_valid, surface_net_radiation_valid,
- # evaporation_valid, total_precipitation_valid, mld_valid, sst_surface_valid, sss_surface_valid,
- # u_surface_valid, v_surface_valid, T_d_valid, S_d_valid, u_d_valid, v_d_valid
- # ), dim=2)
-
-
-
- train_data = torch.cat((sst_surface_train, sss_surface_train,), dim=2) # train_data.shape:torch.Size([5920, 10, 16, 40, 200])
-
- valid_data = torch.cat((sst_surface_valid, sss_surface_valid,), dim=2)
-
- # train_data = sst_surface_train
- # valid_data = sst_surface_valid
-
- # train_size = 1312
- # valid_size = 1760 #前20% 作为验证 剩下的20%的作为测试
-
- sst_train_label = sst_surface_label[14:train_size + 14,:,:]
- sst_valid_label = sst_surface_label[train_size + 14 : valid_size + 14,:,:]
- # sst_test_label = sst_surface_label[valid_size + 14:,:,:]
-
-
- sss_train_label = sss_surface_label[14:train_size + 14,:,:]
- sss_valid_label = sss_surface_label[train_size + 14 : valid_size + 14,:,:]
- # sss_test_label = sss_surface_label[valid_size + 14:,:,:]
-
- sst_train_label = sst_train_label.reshape(-1,1,40,200)
- sst_valid_label = sst_valid_label.reshape(-1,1,40,200)
-
- sss_train_label = sss_train_label.reshape(-1,1,40,200)
- sss_valid_label = sss_valid_label.reshape(-1,1,40,200)
-
- train_label = np.concatenate((sst_train_label, sss_train_label), axis = 1)
- valid_label = np.concatenate((sst_valid_label, sss_valid_label), axis = 1)
- # train_label = sst_train_label
- # valid_label = sst_valid_label
-
- # print('train_label.shape:{}'.format(train_label.shape))
- # print('valid_label.shape:{}'.format(valid_label.shape))
- # print('train_data.shape:{}'.format(train_data.shape))
- # print('valid_data.shape:{}'.format(valid_data.shape))
-
- # print('train_data.shape:{}'.format(train_data.shape)) # train_data.shape:torch.Size([5920, 10, 16, 40, 200])
- # print('valid_data.shape:{}'.format(valid_data.shape)) # valid_data.shape:torch.Size([1952, 10, 16, 40, 200])
- # print('test_data.shape:{}'.format(test_data.shape)) # test_data.shape:torch.Size([1979, 10, 16, 40, 200])
- # print('sst_train_label.shape:{}'.format(sst_train_label.shape)) # sst_train_label.shape:(5920, 14, 40, 200)
- # print('sst_valid_label.shape:{}'.format(sst_valid_label.shape)) # sst_valid_label.shape:(1952, 14, 40, 200)
- # print('sst_test_label.shape:{}'.format(sst_test_label.shape)) # sst_test_label.shape:(1961, 14, 40, 200)
- # print('sss_train_label.shape:{}'.format(sss_train_label.shape)) # sss_train_label.shape:(5920, 14, 40, 200)
- # print('sss_valid_label.shape:{}'.format(sss_valid_label.shape)) # sss_valid_label.shape:(1952, 14, 40, 200)
- # print('sss_test_label.shape:{}'.format(sss_test_label.shape)) # sss_test_label.shape:(1961, 14, 40, 200)
-
-
- #构建数据管道
- class MyDataset(Dataset):
- def __init__(self, data, label):
- self.data = torch.Tensor(data)
- self.label = torch.Tensor(label)
-
- def __len__(self):
- return len(self.label)
-
- def __getitem__(self, idx):
- return self.data[idx], self.label[idx]
-
-
- batch_size1 = 32
- batch_size2 = 32
- batch_size3 = 3000
-
-
- trainset = MyDataset(train_data, train_label)
- trainloader = DataLoader(trainset, batch_size=batch_size1, shuffle=True, drop_last=False,pin_memory=True, num_workers=4)
-
- validset = MyDataset(valid_data, valid_label)
- validloader = DataLoader(validset, batch_size=batch_size2, shuffle=True, drop_last=False,pin_memory=True, num_workers=4)
-
- # testset = MyDataset(test_data, sss_test_label)
- # testloader = DataLoader(testset, batch_size=batch_size3, shuffle=False, drop_last=False,pin_memory=True, num_workers=0)
-
- # print('Qnet_train.shape:{}'.format(sst1_train.shape))
-
-
- model_weights1 = '/model/epo300_lay3_lr0.001_e5_forecastnet_14day_model_weights.pth'
- torch.backends.cudnn.enabled = False
-
- model = AFNONet().cuda()
-
- criterion = nn.MSELoss()
- # 定义优化器
- optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
-
- epochs = 100
- train_losses, valid_losses = [], []
- # best_loss = 2
- best_score = float('inf')
- best_score1 = float('inf')
-
- pred_val= np.zeros((448,2,40,200))
-
- sores = []
- def rmse(y_true, y_preds):
- return np.sqrt(mean_squared_error(y_pred = y_preds, y_true = y_true))
-
-
- for epoch in range(epochs):
- print('Epoch: {}/{}'.format(epoch + 1, epochs))
- # print(var_y)
- #模型训练
- model.train()
- losses = 0
- loss1 = 0
- for i, data in tqdm(enumerate(trainloader)):
- data, label = data
- data = data.cuda()
- label = label.cuda()
-
- # print('data1.shape:{}'.format(data1.shape)) # data1.shape:torch.Size([32, 10, 6, 40, 200]) surface_latent_heat_flux_train surface_sensible_heat_flux_train surface_net_radiation_train evaporation_train total_precipitation_train mld_train
- # print('data2.shape:{}'.format(data2.shape)) # data2.shape:torch.Size([32, 10, 3, 40, 200]) sss_surface_train u_surface_train v_surface_train
- # print('data3.shape:{}'.format(data3.shape)) # data3.shape:torch.Size([32, 10, 40, 200]) S_d_train
- # print('data4.shape:{}'.format(data4.shape)) # data4.shape:torch.Size([32, 10, 2, 40, 200]) u_d_train v_d_train
- # print('label.shape:{}'.format(label.shape)) #label.shape:torch.Size([32, 14, 40, 200])
- # print('data_train.shape:{}'.format(data_train.shape)) # label.shape:torch.Size([32, 14, 40, 200])
- out = model(data)
- # print('out.shape:{}'.format(out.shape))
- # print('label.shape:{}'.format(label.shape))
- # print(out)
- # 偏S/偏t - (E - P) * (S / h) - [u * 偏S/偏x + v * 偏S/偏y ] + H * (w_h + dh/dt * ((S - S_h) / h)) = 0 loss1
- # 偏T/偏t - Q / (p * C_p * h_m) - u * (偏T/偏x) - v * (偏T/偏y) + w_e * ((T - T_d) / h) = 0
-
- # sst_label = label[:,0,:,:]
- # sss_label = label[:,1,:,:]
-
- # sst_out = out[:,0,:,:]
- # sss_out = out[:,1,:,:]
-
- loss = criterion(out, label)
- # loss2 = criterion(sss_out, sss_label)
-
-
- losses += loss
-
- loss.backward()
- optimizer.step()
- train_loss = losses / len(trainloader)
- train_losses.append(train_loss)
-
- print('Training Loss: {:.10f}'.format((train_loss)))
-
- model.eval()
- losses = 0
-
- for i, data in tqdm(enumerate(validloader)):
- data, label = data
- data = data.cuda()
- label = label.cuda()
- optimizer.zero_grad()
-
-
- out = model(data)
- loss = criterion(out, label)
-
- losses += float(loss)
-
- out1 = out.detach().cpu().numpy()
- pred_val[i * batch_size2:(i + 1) * batch_size2] = np.array(out1)
-
- valid_loss = losses / len(validloader)
- valid_losses.append(valid_loss)
-
- valid_label1 = valid_label.reshape(-1,1)
- preds1 = pred_val.reshape(-1,1)
-
- s = rmse(valid_label1,preds1)
- sores.append(s)
- print('Score: {:.3f}'.format(s))
-
- if valid_loss < best_score1: # 求s的最小值 ---》最大值反过来 inf符号也要反过来
- best_score1 = valid_loss
- checkpoint = {'best_score': valid_loss,
- 'state_dict': model.state_dict()}
- torch.save(checkpoint, model_weights1) # if valid_loss < best_loss:
- best_loss = valid_loss
- torch.save(model.state_dict(),
- '/model/pinn_likeICDM_lr0.005_model_300_layer3_2day_e5.pt')
-
- print(sores)
- print(best_score)
- print(s)
|