|
- import numpy as np
- import torch
- from torch import nn
- # from input_data import *
- from sklearn.preprocessing import MinMaxScaler
- print(1)
- # 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
- import torch
- from torch import nn, einsum
- import torch.nn.functional as F
-
- from einops import rearrange, repeat
- from einops.layers.torch import Rearrange
-
- # 判断t是否是元组,如果是,直接返回t;如果不是,则将t复制为元组(t, t)再返回
- # 用来处理当给出的图像尺寸或块尺寸是int类型(如224)时,直接返回为同值元组(如(224, 224))
- def pair(t):
- return t if isinstance(t, tuple) else (t, t)
-
-
- # PreNorm
- class PreNorm(nn.Module):
- def __init__(self, dim, fn):
- super().__init__()
- self.norm = nn.LayerNorm(dim)
- self.fn = fn
- def forward(self, x, **kwargs):
- return self.fn(self.norm(x), **kwargs)
-
-
-
- # Feedforward
-
- class FeedForward(nn.Module):
- def __init__(self, dim, hidden_dim, dropout=0.):
- super().__init__()
- self.net = nn.Sequential(
- nn.Linear(dim, hidden_dim),
- nn.GELU(),
- nn.Dropout(dropout),
- nn.Linear(hidden_dim, dim),
- nn.Dropout(dropout)
- )
- def forward(self, x):
- return self.net(x)
-
- #attention
-
- class Attention(nn.Module):
- def __init__(self, dim, heads=8, dim_head=64, dropout=0.):
- super().__init__()
- inner_dim = dim_head * heads
- project_out = not (heads == 1 and dim_head == dim)
-
- self.heads = heads
- self.scale = dim_head ** -0.5
-
- self.attend = nn.Softmax(dim=-1)
- self.to_qkv = nn.Linear(dim, inner_dim * 3, bias=False)
-
- self.to_out = nn.Sequential(
- nn.Linear(inner_dim, dim),
- nn.Dropout(dropout),
- ) if project_out else nn.Identity()
-
- def forward(self, x):
- b, n, _, h = *x.shape, self.heads
- qkv = self.to_qkv(x).chunk(3, dim=-1) # (b, n(65), dim*3) ---> 3 * (b, n, dim)
- q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h=h), qkv) # q, k, v (b, h, n, dim_head(64))
-
- dots = einsum('b h i d, b h j d -> b h i j', q, k) * self.scale
-
- attn = self.attend(dots)
-
- out = einsum('b h i j, b h j d -> b h i d', attn, v)
- out = rearrange(out, 'b h n d -> b n (h d)')
- return self.to_out(out)
-
-
- # Transformet
-
- class Transformer(nn.Module):
- def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout=0.):
- super().__init__()
- self.layers = nn.ModuleList([])
- for _ in range(depth):
- self.layers.append(nn.ModuleList([
- PreNorm(dim, Attention(dim, heads=heads, dim_head=dim_head, dropout=dropout)),
- PreNorm(dim, FeedForward(dim, mlp_dim, dropout=dropout))
- ]))
-
- def forward(self, x):
- for attn, ff in self.layers:
- x = attn(x) + x
- x = ff(x) + x
- return x
-
-
- #VIT
- class ViT(nn.Module):
- def __init__(self, *, image_size, patch_size, num_classes, dim, depth, heads, mlp_dim, pool='cls', channels=3, dim_head=64, dropout=0., emb_dropout=0.
- ):
- super().__init__()
- image_height, image_width = pair(image_size) # 8 8
- patch_height, patch_width = pair(patch_size) # 32 32
-
- assert image_height % patch_height ==0 and image_width % patch_width == 0
-
- num_patches = (image_height // patch_height) * (image_width // patch_width)
- patch_dim = channels * patch_height * patch_width
- assert pool in {'cls', 'mean'}
-
- # self.to_patch_embedding = nn.Sequential(
- # Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1=patch_height, p2=patch_width),
- # nn.Linear(patch_dim, dim)
- # )
-
- self.pos_embedding = nn.Parameter(torch.randn(1, num_patches+1, dim))
- self.cls_token = nn.Parameter(torch.randn(1, 1, dim)) # nn.Parameter()定义可学习参数
- self.dropout = nn.Dropout(emb_dropout)
-
- self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim, dropout)
-
- self.pool = pool
- self.to_latent = nn.Identity()
-
- self.mlp_head = nn.Sequential(
- nn.LayerNorm(dim),
- nn.Linear(dim, num_classes)
- )
- self.lst = nn.LayerNorm(dim)
-
- self.conv1 = nn.Conv3d(in_channels=6, out_channels=1, kernel_size=(3,3,3), padding=(1,1,1))
- self.conv2 = nn.Conv3d(in_channels=10, out_channels=1, kernel_size=(3,3,3), padding=(1,1,1))
- self.conv = nn.Conv2d(in_channels=41, out_channels=40, kernel_size=(3,3), padding=(1,1))
-
- def forward(self, x):
- b, n, t, c, _ = x.shape # b表示batchSize, n表示每个块的空间分辨率, _表示一个块内有多少个值 #batch=16 n=64 _=1024
-
- x = self.conv1(x)
-
- N,B,T,C,H = x.size()
- x = x.permute(0,2,1,3,4)
-
- x = self.conv2(x)
-
- x = x.reshape(-1,40,200)
-
- b, n, _ = x.shape
-
- cls_tokens = repeat(self.cls_token, '() n d -> b n d', b=b) # self.cls_token: (1, 1, dim) -> cls_tokens: (batchSize, 1, dim)
- x = torch.cat((cls_tokens, x), dim=1)
- x += self.pos_embedding[:, :(n+1)] # 加位置嵌入(直接加) (b, 65, dim)
- x = self.dropout(x)
- x = self.transformer(x) # (b, 65, dim)
- x = self.to_latent(x) # Identity (b, dim)
-
- x = x.unsqueeze(-1)
-
- x_lst = self.conv(x)
-
- return x_lst
-
-
- print(2)
- for seed in range(2023,2024):
- for date_append in range(0, 15):
- 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)
-
- print(3)
- data = np.load(r'/dataset/10day_for_14day_all_variables_deep1_pacific_10_19_SSS_learn.npz')
- print(data.files)
- # 2000-01-01 ---> 2019-12-31
- evaporation = data['evaporation'][:]
- total_precipitation = data['total_precipitation'][:]
- mld = data['mld'][:]
- # sst_surface = data['sst_surface'][:]
- sss_surface = data['sst_deep'][:]
- u_surface = data['u_deep'][:]
- v_surface = data['v_deep'][:]
- sss_surface_label = data['sst_deep_label'][:]
- # print(sst_surface)
-
- # print(sst_surface.shape)
- print(sss_surface.shape) #(3642, 10, 40, 200)
-
- print(sss_surface_label.shape) #(3638, 40, 200)
-
- dS_dt = data['dT_dt'][:]
- dS_dx = data['dT_dx'][:]
- dS_dy = data['dT_dy'][:]
-
-
-
- train_size = 2208
- valid_size = 2912 # 前20% 作为验证 剩下的20%的作为测试
- # surface_latent_heat_flux = surface_latent_heat_flux + surface_sensible_heat_flux + surface_net_radiation
-
-
- evaporation = evaporation.reshape(-1, 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, 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, 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, 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, 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, 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, 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:,:,:,:,:]
-
- dS_dt = dS_dt.transpose(1,0,2,3)
- dS_dt = dS_dt.reshape(-1, 1, 10, 40, 200)
- dS_dt = torch.Tensor(dS_dt)
- dS_dt_train = dS_dt[0:train_size, :, :, :]
- dS_dt_valid = dS_dt[train_size:valid_size, :, :, :]
-
- dS_dx = dS_dx.transpose(1,2,0,3)
- dS_dx = dS_dx.reshape(-1, 1, 10, 40, 200)
- dS_dx = torch.Tensor(dS_dx)
- dS_dx_train = dS_dx[0:train_size, :, :, :]
- dS_dx_valid = dS_dx[train_size:valid_size, :, :, :]
-
- dS_dy = dS_dy.transpose(1,2,3,0)
- dS_dy = dS_dy.reshape(-1, 1, 10, 40, 200)
- dS_dy = torch.Tensor(dS_dy)
- dS_dy_train = dS_dy[0:train_size, :, :, :]
- dS_dy_valid = dS_dy[train_size:valid_size, :, :, :]
-
-
- train_data = torch.cat((evaporation_train, total_precipitation_train, sss_surface_train, mld_train, u_surface_train, v_surface_train,
- dS_dt_train, dS_dx_train, dS_dy_train
- ), dim=1) # train_data.shape:torch.Size([5920, 10, 16, 40, 200])
-
- valid_data = torch.cat((evaporation_valid, total_precipitation_valid, sss_surface_valid, mld_valid, u_surface_valid, v_surface_valid,
- dS_dt_valid, dS_dx_valid, dS_dy_valid
- ), dim=1)
-
-
-
- print(train_data.shape)
- print(valid_data.shape)
-
- sss_train_label = sss_surface_label[10 :train_size + 10 , date_append, :, :]
- sss_valid_label = sss_surface_label[train_size + 10: valid_size + 10 , date_append, :, :]
-
-
- sss_train_label = sss_train_label.reshape(-1, 1, 40, 200)
- sss_valid_label = sss_valid_label.reshape(-1, 1, 40, 200)
-
- train_label = sss_train_label
- valid_label = sss_valid_label
- print(train_label.shape)
-
-
-
- # 构建数据管道
- 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)
-
- model_weights1 = '/model/epo200_lr0.001_e{}_VIT_2D_{}day_model_weights_SST_deep8.pth'.format(seed, date_append)
-
- model = ViT(
- image_size = 256,
- patch_size = 32,
- num_classes = 1000,
- dim = 200,
- depth = 8, #1
- heads = 12, #2
- mlp_dim = 2048,
- dropout = 0.1,
- emb_dropout = 0.1
- ).cuda()
-
- model = model.cuda()
-
- criterion = nn.MSELoss()
- # 定义优化器
- optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
-
- epochs = 200
- train_losses, valid_losses = [], []
- # best_loss = 2
- best_score = float('inf')
- best_score1 = float('inf')
-
- pred_val = np.zeros((704, 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 data, label in tqdm(trainloader):
- data = data.cuda()
- label = label.cuda()
- optimizer.zero_grad()
-
- # B,T,C,H,W = data.size()
- # data1 = data.reshape(B,T*C,H*W)
- data_train = data[:,:6,:,:,:]
- # print('data.shape:{}'.format(data.shape)) #data1.shape:torch.Size([32, 60, 8000])
- out = model(data_train)
- # print(out)
- # print('out.shape:{}'.format(out.shape)) #out.shape:torch.Size([32, 8000])
-
- 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
- with torch.no_grad():
- for i, data in tqdm(enumerate(validloader)):
- data, label = data
- data = data.cuda()
- label = label.cuda()
- optimizer.zero_grad()
- # B,T,C,H,W = data.size()
- # data1 = data.reshape(B,T*C,H*W)
- data_train = data[:,:6,:,:,:]
- out = model(data_train)
- out = out.reshape(-1, 1, 40, 200)
- 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/VIT_2D_lr0.001_model_200_{}day_e{}_SST_deep8.pt'.format(date_append, seed))
-
- print(sores)
- print(best_score)
- print(s)
|