|
- import torch
- from torch import nn
- from sklearn.preprocessing import MinMaxScaler
- from sklearn.metrics import mean_squared_error
- import random
- import numpy as np
- import torch
- from tqdm import tqdm
- from torch.utils.data import Dataset, DataLoader
- mm = MinMaxScaler()
-
-
-
- class ConvNet(nn.Module):
- def __init__(self,input_dim, input_dim1, hidden_dim, kernel_size1,padding1,kernel_size2,padding2,kernel_size3,padding3):
- super(ConvNet, self).__init__()
- self.input_dim = input_dim
- self.input_dim1 = input_dim1
-
- self.hidden_dim = hidden_dim
- self.kernel_size1 = kernel_size1
- self.padding1 = padding1
-
- self.kernel_size2 = kernel_size2
- self.kernel_size3 = kernel_size3
- self.padding2 = padding2
- self.padding3 = padding3
-
- self.gelu = nn.GELU()
-
- #stage1
- self.conv1 = nn.Sequential(nn.Conv3d(in_channels=self.input_dim ,
- out_channels=16,
- kernel_size=kernel_size1,
- padding=self.padding1,
- ))
- self.conv2 = nn.Sequential(nn.Conv3d(in_channels=16,
- out_channels=32,
- kernel_size=kernel_size2,
- padding=self.padding2,
- ))
- self.conv3 = nn.Sequential(nn.Conv3d(in_channels=32,
- out_channels=7,
- kernel_size=kernel_size3,
- padding=self.padding3,
- ))
- # self.conv4 = nn.Sequential(nn.Conv2d(in_channels=64,
- # out_channels=1,
- # kernel_size=kernel_size,
- # padding=self.padding,
- # ))
- # stage2
- self.conv11 = nn.Sequential(nn.Conv3d(in_channels=self.input_dim1,
- out_channels=16,
- kernel_size=kernel_size1,
- padding=self.padding1,
- ))
- self.conv22 = nn.Sequential(nn.Conv3d(in_channels=16,
- out_channels=32,
- kernel_size=kernel_size2,
- padding=self.padding2,
- ))
- self.conv33 = nn.Sequential(nn.Conv3d(in_channels=32,
- out_channels=10,
- kernel_size=kernel_size3,
- padding=self.padding3,
- ))
- # self.fc = nn.Linear
-
- # 前馈网络过程
- def forward(self, x): #x 为(batch, channel, h, w)
-
- out = self.conv1(x) #x 为(batch, channel, h, w)
- out = self.gelu(out)
- out = self.conv2(out)
- out = self.gelu(out)
- out = self.conv3(out)
- out = self.gelu(out)
-
- print('out.shape:{}'.format(out.shape))
- out1 = out.permute(0,2,1,3,4)
- out1 = self.conv11(out1)
- out1 = self.gelu(out1)
- out1 = self.conv22(out1)
- out1 = self.gelu(out1)
- out1 = self.conv33(out1)
- out1 = self.gelu(out1)
-
- out11 = out1.permute(0,2,1,3,4)
- # out = self.conv4(out)
- return out11
-
- for iii in range(1,6):
- for epoch111 in range(100,600,100):
-
- 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(iii)
-
- data = np.load('/tmp/dataset/7day_for_Nday_data_openi_09_17_atlantic_area_last.npz')
- print(data.files) # ['hycom_temp', 'slfh', 'sshf', 'ssr', 'str', 'mld', 'analysis_temp', 'u', 'v', 'T_d', 'u_d', 'v_d', 'xx', 'yy']
-
- hycom_temp = data['hycom_temp'][:] # (3281, 7, 7, 41, 201)
- slfh = data['slfh'][:] # (3281, 7, 41, 201)
- sshf = data['sshf'][:] # (3281, 7, 41, 201)
- ssr = data['ssr'][:] # (3281, 7, 41, 201)
- str = data['str'][:] # (3281, 7, 41, 201)
- mld = data['mld'][:] # (3281, 7, 41, 201)
- analysis_temp = data['analysis_temp'][:] # (3281, 7, 7, 41, 201)
- u = data['u'][:] # (3281, 7, 7, 41, 201)
- v = data['v'][:] # (3281, 7, 7, 41, 201)
- T_d = data['T_d'][:] # (3281, 7, 41, 201)
- u_d = data['u_d'][:] # (3281, 7, 41, 201)
- v_d = data['v_d'][:] # (3281, 7, 41, 201)
- xx = data['xx'][:] # (3281, 7, 41, 201)
- yy = data['yy'][:] # (3281, 7, 41, 201)
-
-
- analysis_temp = analysis_temp.transpose(0,2,1,3,4)
-
- # print(analysis_temp.shape)
-
-
- train_size = 1952 # 前60%
- valid_size = 2624 # 中间20% 作为验证 剩下的20%的作为测试
-
- Q_net = (slfh + sshf + ssr + str) / 86400
- # print(Q_net)
- hycom_temp = hycom_temp.transpose(0,2,1,3,4)
- hycom_temp = hycom_temp.reshape(-1,7,7,41,81)
- hycom_temp = torch.Tensor(hycom_temp)
- hycom_temp_train = hycom_temp[0:train_size,:,:,:,:]
- hycom_temp_valid = hycom_temp[train_size:valid_size,:,:,:,:]
- hycom_temp_test = hycom_temp[valid_size:,:,:,:,:]
-
-
- # slfh = slfh.reshape(-1,1,1,7,41,201)
- # slfh = torch.Tensor(slfh)
- # slfh_train = slfh[0:train_size,:,:,:,:]
- # slfh_valid = slfh[train_size:valid_size,:,:,:,:]
- # slfh_test = slfh[valid_size:,:,:,:,:]
- #
- # sshf = sshf.reshape(-1,1,1,7,41,201)
- # sshf = torch.Tensor(sshf)
- # sshf_train = sshf[0:train_size,:,:,:,:]
- # sshf_valid = sshf[train_size:valid_size,:,:,:,:]
- # sshf_test = sshf[valid_size:,:,:,:,:]
- #
- #
- # ssr = ssr.reshape(-1,1,1,7,41,201)
- # ssr = torch.Tensor(ssr)
- # ssr_train = ssr[0:train_size,:,:,:,:]
- # ssr_valid = ssr[train_size:valid_size,:,:,:,:]
- # ssr_test = ssr[valid_size:,:,:,:,:]
- #
- #
- # str = str.reshape(-1,1,1,7,41,201)
- # str = torch.Tensor(str)
- # str_train = str[0:train_size,:,:,:,:]
- # str_valid = str[train_size:valid_size,:,:,:,:]
- # str_test = str[valid_size:,:,:,:,:]
-
- Q_net = Q_net.reshape(-1,1,7,41,81)
- 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:,:,:,:,:]
-
-
- mld = mld.reshape(-1,1,7,41,81)
- mld = torch.Tensor(mld)
- mld_train = mld[0:train_size,:,:,:,:]
- mld_valid = mld[train_size:valid_size,:,:,:,:]
- mld_test = mld[valid_size:,:,:,:,:]
-
-
- analysis_temp = analysis_temp.reshape(-1,7,10,41,81)
- analysis_temp = torch.Tensor(analysis_temp)
- # analysis_temp_train = analysis_temp[0:train_size,:,:,:,:]
- # analysis_temp_valid = analysis_temp[train_size:valid_size,:,:,:,:]
- # analysis_temp_test = analysis_temp[valid_size:,:,:,:,:]
-
- u = u.transpose(0,2,1,3,4)
- u = u.reshape(-1,7,7,41,81)
- u = torch.Tensor(u)
- u_train = u[0:train_size,:,:,:,:]
- u_valid = u[train_size:valid_size,:,:,:,:]
- u_test = u[valid_size:,:,:,:,:]
-
- v = v.transpose(0,2,1,3,4)
- v = v.reshape(-1,7,7,41,81)
- v = torch.Tensor(v)
- v_train = v[0:train_size,:,:,:,:]
- v_valid = v[train_size:valid_size,:,:,:,:]
- v_test = v[valid_size:,:,:,:,:]
-
-
- T_d = T_d.reshape(-1,1,7,41,81)
- 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:,:,:,:,:]
-
-
- u_d = u_d.reshape(-1,1,7,41,81)
- 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,1,7,41,81)
- 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,7,41,81)
- 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,7,41,81)
- yy = torch.Tensor(yy)
- yy_train = yy[0:train_size,:,:,:,:]
- yy_valid = yy[train_size:valid_size,:,:,:,:]
- yy_test = yy[valid_size:,:,:,:,:]
-
- print('hycom_temp_train.shape:{}'.format(hycom_temp_train.shape))
- print('Q_net_train.shape:{}'.format(Q_net_train.shape))
- print('mld_train.shape:{}'.format(mld_train.shape))
- print('u_train.shape:{}'.format(u_train.shape))
- print('v_train.shape:{}'.format(v_train.shape))
- print('T_d_train.shape:{}'.format(T_d_train.shape))
- print('v_d_train.shape:{}'.format(v_d_train.shape))
-
-
- train_data = torch.cat((hycom_temp_train, Q_net_train, mld_train, u_train, v_train, T_d_train, v_d_train), dim=1) # train_data.shape:torch.Size([1952, 7, 32, 41, 81]) # 数据 第一维时间 第二维深度和变量 第三维7dat seq 4 5lat lon
-
-
- print('train_data.shape:{}'.format(train_data.shape))
-
- valid_data = torch.cat((hycom_temp_valid, Q_net_valid, mld_valid, u_valid, v_valid, T_d_valid, v_d_valid), dim=1) # valid_data.shape:torch.Size([672, 7, 32, 41, 81])
- print('valid_data.shape:{}'.format(valid_data.shape)) #([326, 3, 12, 15])
-
- test_data = torch.cat((hycom_temp_test, Q_net_test, mld_test, u_test, v_test, T_d_test, v_d_test), dim=1) #
- print('test_data.shape:{}'.format(test_data.shape)) #([326, 3, 12, 15])
-
- train_label = analysis_temp[7:train_size + 7,:,:,:] # train_label.shape:torch.Size([1952, 7, 7, 41, 81])
-
- valid_label = analysis_temp[train_size + 7:valid_size + 7,:,:,:] # valid_label.shape:torch.Size([672, 7, 7, 41, 81])
-
-
- test_label = analysis_temp[valid_size + 7: 3646,:,:,:] # test_label.shape:torch.Size([650, 7, 7, 41, 81])
-
-
- print('test_label.shape:{}'.format(test_label.shape)) # test_label.shape:(63, 1, 6, 27)
-
- train_label = torch.Tensor(train_label)
- valid_label = torch.Tensor(valid_label)
- test_label = torch.Tensor(test_label)
- print('train_label.shape:{}'.format(train_label.shape)) #train_label.shape:torch.Size([224, 1, 6, 27])
- test_label11 = test_label
- valid_label11 = valid_label
-
- #构建数据管道
- 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(input_data, sst)
- # trainset = MyDataset(train_data,train_label)
- # trainloader = DataLoader(trainset, batch_size=batch_size, shuffle=True)
- # batch_size = 1826
- 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, test_label)
- testloader = DataLoader(testset, batch_size=batch_size3, shuffle=False, drop_last=False,pin_memory=True, num_workers=4)
-
- # print('Qnet_train.shape:{}'.format(sst1_train.shape))
-
-
-
- print('-----------------------train_black_cnn----------------------------')
-
- # model_weights = './300e1_layer2_black_blck_8day_model_weights.pth'
- model_weights1 = '/cnn_black_WM_gelu_epo{}_lay1_e{}_black_ahead_1_day_model_weights.pth'.format(epoch111,iii)
-
- model = ConvNet(input_dim=25,input_dim1=7,hidden_dim=1,kernel_size1=(1,3,3),padding1=(0,1,1),kernel_size2=(1,3,3), padding2=(0,1,1), kernel_size3=(1,3,3), padding3=(0,1,1)).cuda()
- # model = ConvNet(input_dim=9,hidden_dim=1,kernel_size1=(3,3,3),padding1=(1,1,1)).cuda()
- # out = model(a) #torch.Size([1826, 12, 15])
- # print(out.shape)
- criterion = nn.MSELoss()
- # 定义优化器
- optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
-
- epochs = 400
- train_losses, valid_losses = [], []
- # best_loss = 2
- best_score = float('inf')
- best_score2 = float('inf')
-
- preds = np.zeros((672,7,10,41,81))
- # preds = np.expand_dims(preds, axis=1)#preds.reshape(100,1,12,15)
-
- 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(epoch111):
- print('Epoch: {}/{}'.format(epoch + 1, epoch111))
- # print(var_y)
- #模型训练
- model.train()
- losses = 0
- for data, label in tqdm(trainloader):
- data = data.cuda()
- print('data.shape:{}'.format(data.shape)) #data.shape:torch.Size([32, 9, 3, 6, 27])
- label = label.cuda()
- optimizer.zero_grad()
- out = model(data) #
- print('out.shape:{}'.format(out.shape)) # out.shape:torch.Size([32, 1, 3, 6, 27])
- print('label.shape:{}'.format(label.shape))# label.shape:torch.Size([32, 1, 6, 27])
- loss = criterion(out, label)
- # losses +=loss.cpu().detach().numpy
- losses += loss
-
-
- # 反向传播
- loss.backward()
- optimizer.step()
- train_loss = losses / len(trainloader)
- train_losses.append(train_loss)
- print('Training Loss: {:.3f}'.format(train_loss))
-
- # # 模型验证
- # model.eval()
- losses = 0
- with torch.no_grad():
- # for i in range(1):
- for i, data in tqdm(enumerate(validloader)):
- # print('i:{}'.format(i))
- data, label = data
- data = data.cuda()
- label = label.cuda()
- optimizer.zero_grad()
- pred = model(data)
-
- print('pred_shape:{}'.format(pred.shape)) #pred_shape:torch.Size([32, 1, 3, 6, 27])
-
- print('label_shape:{}'.format(label.shape)) # label_shape:torch.Size([32, 1, 6, 27])
-
- loss = criterion(pred, label)
- losses += loss
- # pred = pred.reshape(-1,1)
- # preds[i * batch_size2:(i + 1) * batch_size2] = Nino_pred.cpu()
- # print(preds.shape)
- # preds[i * batch_size2:(i + 1) * batch_size2] = np.array(tmp)
- preds[i * batch_size2:(i + 1) * batch_size2] = pred.cpu()
-
- valid_loss = losses / len(validloader)
- valid_losses.append(valid_loss)
- # print('Validation Loss: {:.3f}'.format(valid_loss))
-
- print('valid_label_shape:{}'.format(valid_label.shape)) # valid_label_shape:torch.Size([352, 1, 6, 27])
- valid_label1 = valid_label
- valid_label2 = valid_label1.reshape(-1,1)
- preds1 = preds.reshape(-1,1)
- print('valid_label1.shape:{}'.format(valid_label2.shape))#valid_label1.shape:torch.Size([57024, 1])
- print('preds1.shape:{}'.format(preds1.shape)) #preds1.shape:torch.Size([15552, 1])0
- s = rmse(valid_label2,preds1)
- sores.append(s)
- print('Score: {:.3f}'.format(s))
- if valid_loss < best_score2: # 求s的最小值 ---》最大值反过来 inf符号也要反过来
- best_score2 = 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(),
- '/cnn_black_WM_gelu_e{}_layer1_ahead_1_day_e{}.pt'.format(epoch111,iii))
- # print(s)
- # print(s)
- print(sores)
- print(best_score)
- print(s)
|