|
- import numpy as np
- import torch
- from torch import nn
- # from input_data import *
- from sklearn.preprocessing import MinMaxScaler
-
- # ConvLSTM Cell
- import random
-
- import numpy as np
- import torch
- from tqdm import tqdm
- from sklearn.metrics import mean_squared_error
- from torch import nn
- from torch.utils.data import Dataset, DataLoader
-
-
- mm = MinMaxScaler()
-
- #3dcnn 的输入(N, C, D, H ,W) ---> (N, C1, D1, H1, W1)
- #strid = 1 ---> kernel_size = input_dim + 1 - output_dim + 2*padding
-
- class unetConv2(nn.Module):
- def __init__(self, in_size, out_size, is_batchnorm, n=2, ks=3, stride=1, padding=1):
- super(unetConv2, self).__init__()
- self.n = n
- self.ks = ks
- self.stride = stride
- self.padding = padding
- s = stride
- p = padding
- if is_batchnorm:
- for i in range(1, n + 1):
- conv = nn.Sequential(nn.Conv2d(in_size, out_size, ks, s, p),
- nn.BatchNorm2d(out_size),
- nn.ReLU(inplace=True), )
- setattr(self, 'conv%d' % i, conv)
- in_size = out_size
-
- else:
- for i in range(1, n + 1):
- conv = nn.Sequential(nn.Conv2d(in_size, out_size, ks, s, p),
- nn.ReLU(inplace=True), )
- setattr(self, 'conv%d' % i, conv)
- in_size = out_size
-
-
- def forward(self, inputs):
- x = inputs
- for i in range(1, self.n + 1):
- conv = getattr(self, 'conv%d' % i)
- x = conv(x)
-
- return x
-
-
- class unetUp(nn.Module):
- def __init__(self, in_size, out_size, is_deconv, n_concat=2):
- super(unetUp, self).__init__()
- self.conv = unetConv2(in_size + (n_concat - 2) * out_size, out_size, False)
- if is_deconv:
- self.up = nn.ConvTranspose2d(in_size, out_size, kernel_size=2, stride=2, padding=0)
- else:
- self.up = nn.Sequential(
- nn.UpsamplingBilinear2d(scale_factor=2),
- nn.Conv2d(in_size, out_size, 1))
-
- def forward(self, high_feature, *low_feature):
- outputs0 = self.up(high_feature)
- for feature in low_feature:
- outputs0 = torch.cat([outputs0, feature], 1)
- return self.conv(outputs0)
-
-
- class UNet_learn(nn.Module):
-
- def __init__(self, in_channels=60, n_classes=1, feature_scale=2, is_deconv=True, is_batchnorm=True):
- super(UNet_learn, self).__init__()
- self.in_channels = in_channels
- self.feature_scale = feature_scale
- self.is_deconv = is_deconv
- self.is_batchnorm = is_batchnorm
-
- filters = [64, 128, 256, 512]
- filters = [int(x / self.feature_scale) for x in filters]
-
- # downsampling
- self.maxpool = nn.MaxPool2d(kernel_size=2)
- self.conv1 = unetConv2(self.in_channels, filters[0], self.is_batchnorm)
- self.conv2 = unetConv2(filters[0], filters[1], self.is_batchnorm)
- self.conv3 = unetConv2(filters[1], filters[2], self.is_batchnorm)
- # self.conv4 = unetConv2(filters[2], filters[3], self.is_batchnorm)
- self.center = unetConv2(filters[2], filters[3], self.is_batchnorm)
- # upsampling
- # self.up_concat4 = unetUp(filters[4], filters[3], self.is_deconv)
- self.up_concat3 = unetUp(filters[3], filters[2], self.is_deconv)
- self.up_concat2 = unetUp(filters[2], filters[1], self.is_deconv)
- self.up_concat1 = unetUp(filters[1], filters[0], self.is_deconv)
- # final conv (without any concat)
- self.final = nn.Conv2d(filters[0], n_classes, 1)
-
-
- def forward(self, inputs):
- # print('input.shape:{}'.format(inputs.shape)) # 输入([2, 1, 64, 64]) | ([2, 10, 40, 200])
- conv1 = self.conv1(inputs) # 16*512*512
- # print('conv1.shape:{}'.format(conv1.shape)) # 输入([2, 1, 64, 64]) ---> ([2, 32, 64, 64]) | ([2, 10, 40, 200]) ---> ([2, 32, 40, 200])
- maxpool1 = self.maxpool(conv1) # 16*256*256
- # print('maxpool1.shape:{}'.format(maxpool1.shape)) # 输入([2, 32, 64, 64]) ---> ([2, 32, 32, 32]) | ([2, 32, 40, 200]) ---> ([2, 32, 20, 100])
- conv2 = self.conv2(maxpool1) # 32*256*256
- # print('conv2.shape:{}'.format(conv2.shape)) # 输入([2, 32, 32, 32]) ---> ([2, 64, 32, 32]) | ([2, 32, 20, 100]) ---> ([2, 64, 20, 100])
- maxpool2 = self.maxpool(conv2) # 32*128*128
- # print('maxpool2.shape:{}'.format(maxpool2.shape)) # 输入([2, 64, 32, 32]) ---> ([2, 64, 16, 16]) | ([2, 64, 20, 100]) ---> ([2, 64, 10, 50])
- conv3 = self.conv3(maxpool2) # 64*128*128
- # print('conv3.shape:{}'.format(conv3.shape)) # 输入([2, 64, 16, 16]) ---> ([2, 128, 16, 16]) | ([2, 64, 10, 50]) ---> ([2, 128, 10, 50])
- maxpool3 = self.maxpool(conv3) # 64*64*64
- # print('maxpool3.shape:{}'.format(maxpool3.shape)) # 输入([2, 128, 16, 16]) ---> ([2, 128, 8, 8]) | ([2, 128, 10, 50]) --->([2, 128, 5, 25])
- # conv4 = self.conv4(maxpool3) # 128*64*64
- # # print('conv4.shape:{}'.format(conv4.shape)) # 输入([2, 128, 8, 8])---> ([2, 256, 8, 8])
- # maxpool4 = self.maxpool(conv4) # 128*32*32
- # # print('maxpool4.shape:{}'.format(maxpool4.shape)) # 输入([2, 256, 8, 8])--> ([2, 256, 4, 4])
-
- center = self.center(maxpool3) # 256*32*32
- # print('center.shape:{}'.format(center.shape)) # 输入([2, 256, 4, 4])--> ([2, 512, 4, 4]) | ([2, 128, 5, 25]) ---> ([2, 128, 5, 25])
- # up4 = self.up_concat4(center, conv3) # 128*64*64
- # print('up4.shape:{}'.format(up4.shape)) # 输入([2, 512, 4, 4])--> ([2, 256, 8, 8])
- up3 = self.up_concat3(center, conv3) # 64*128*128
- # print('up3.shape:{}'.format(up3.shape)) # 输入([2, 256, 8, 8])--> ([2, 128, 16, 16])
- up2 = self.up_concat2(up3, conv2) # 32*256*256
- # print('up2.shape:{}'.format(up2.shape)) # 输入([2, 128, 16, 16])-> ([2, 64, 32, 32])
- up1 = self.up_concat1(up2, conv1) # 16*512*512
- # print('up1.shape:{}'.format(up1.shape)) # 输入([2, 64, 32, 32])-> ([2, 32, 64, 64])
-
- final = self.final(up1)
-
- return final
-
-
-
- for seed in range(1, 2):
- for date_append in range(1, 2):
- 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(seed)
-
- # 需要 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, :, :]
-
- print(sst_surface)
- print(np.max(sst_surface))
- print(np.min(sst_surface))
-
- print(sst_surface.shape)
- print(sss_surface.shape)
-
- # 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_radiation = 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)
-
- Q_net = surface_latent_heat_flux + surface_sensible_heat_flux + surface_net_radiation
-
- Q_net = Q_net/86400
-
-
- train_size = 1312
- valid_size = 1760 # 前20% 作为验证 剩下的20%的作为测试
- # surface_latent_heat_flux = surface_latent_heat_flux + surface_sensible_heat_flux + surface_net_radiation
-
- Q_net = Q_net.reshape(-1,10, 40, 200)
- Q_net = torch.Tensor(Q_net)
- Q_net_train = Q_net[0:train_size, :, :, :]
- Q_net_valid = Q_net[train_size:valid_size, :, :, :]
- # Q_net_test = Q_net[valid_size:,:,:,:,:]
-
- surface_latent_heat_flux = surface_latent_heat_flux.reshape(-1, 10, 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, 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, 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, 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, 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, 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, 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, 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, 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, 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, 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, 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, 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, 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, 1, 10, 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, 1, 10, 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((Q_net_train, mld_train, sst_surface_train, u_surface_train, v_surface_train, T_d_train
- ), dim=1) # train_data.shape:torch.Size([5920, 10, 16, 40, 200])
-
- valid_data = torch.cat((Q_net_valid,mld_valid, sst_surface_valid,u_surface_valid, v_surface_valid, T_d_valid
- ), dim=1)
-
- # print(train_data.shape)
-
-
- # train_data = sst_surface_train
- # # train_data.shape:torch.Size([5920, 10, 16, 40, 200])
-
- # valid_data = sst_surface_valid
-
- print(sst_surface_train)
-
-
- # train_data = torch.cat((sst_surface_train, sss_surface_train,), dim=1) # train_data.shape:torch.Size([5920, 10, 16, 40, 200])
-
- # valid_data = torch.cat((sst_surface_valid, sss_surface_valid,), dim=1)
-
- print(train_data.shape)
- print(valid_data.shape)
-
- sst_train_label = sst_surface_label[10 + date_append: train_size + 10 + date_append, :, :]
- sst_valid_label = sst_surface_label[train_size + 10 + date_append: valid_size + 10 + date_append, :, :]
- # sst_test_label = sst_surface_label[valid_size + 14:,:,:]
-
- sss_train_label = sss_surface_label[10 + date_append:train_size + 10 + date_append, :, :]
- sss_valid_label = sss_surface_label[train_size + 10 + date_append: valid_size + 10 + date_append, :, :]
- # sss_test_label = sss_surface_label[valid_size + 10 + date_append:,:,:]
-
- 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
-
- # 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 = 16
- batch_size2 = 16
- 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)
-
- model_weights1 = '/model/epo200_lay4_lr0.001_e{}_Unet_2D_{}day_model_weights_SST_relu.pth'.format(seed, date_append)
-
- model = UNet_learn().cuda()
-
- criterion = nn.MSELoss()
- # 定义优化器
- optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
-
- epochs = 200
- train_losses, valid_losses = [], []
- # best_loss = 2
- best_score1 = float('inf')
-
- pred_val = np.zeros((448, 1, 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.shape)) # data_train.shape:torch.Size([32, 12, 10, 40, 200])
- # print('data.shape:{}'.format(data.shape)) # data.shape:torch.Size([16, 6, 10, 40, 200])
- # data = data.permute(0,2,1,3,4)
- out = model(data)
- # print(out)
- # 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,:,:]
- out = out.reshape(-1, 1, 40, 200)
- 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()
- # data = data.permute(0,2,1,3,4)
- out = model(data)
- out = out.reshape(-1, 1, 40, 200)
- loss = criterion(out, label)
-
- loss1 += float(loss)
-
- out1 = out.detach().cpu().numpy()
- pred_val[i * batch_size2:(i + 1) * batch_size2] = np.array(out1)
-
- valid_loss = loss1 / 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/Unet_2D_lr0.001_model_200_layer4_{}day_e{}_SST_relu.pt'.format(date_append, seed))
-
- print(sores)
- print(best_score)
- print(s)
|