File size: 3,193 Bytes
2b21abc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 |
import os
import torch
import torch.nn.functional as F
from torch import nn, Tensor
import numpy as np
import h5py
from torch.utils.data import DataLoader, Dataset
from torch.utils.data import Subset
from sklearn.model_selection import train_test_split
import torch.optim as optim
from model_convlstm import ionexDataset, train_npy, nstepsin, nstepsout, stride, EncoderDecoderConvLSTM, max_epochs
from model_LARRES import larres
# ionexData = ionexDataset(train_npy, nstepsin=nstepsin, nstepsout=nstepsout, stride=stride)
# train_data, val_data = ionexData.split_train_val(val_split=0.2)
#
# train_loader = DataLoader(train_data, batch_size=16, num_workers=0)
# val_loader = DataLoader(val_data, batch_size=16, num_workers=0)
ionexData = ionexDataset(train_npy, nstepsin=nstepsin, nstepsout=nstepsout, stride=stride)
train_data, val_data = ionexData.split_train_val(val_split=0.2)
train_loader = DataLoader(train_data, batch_size=16, num_workers=0)
val_loader = DataLoader(val_data, batch_size=16, num_workers=0)
for X, y in train_loader:
print(f"Shape of X: {X.shape} {X.dtype} [N, C, H, W]")
print(f"Shape of Y: {y.shape} {y.dtype}")
break
print(f"Training samples: {len(train_loader.dataset)}")
# print(f"Validation samples: {len(val_loader.dataset)}")
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model=larres().to(device)
# model.load_state_dict(torch.load("best_model.pth"))
optimizer = optim.Adam(model.parameters(), lr=0.0001)
criterion = nn.L1Loss()
# 训练和验证
best_val_loss = float('inf')
# num_epochs = 50
for epoch in range(max_epochs):
# 训练阶段
model.train()
all_loss = 0
for batch_idx, (data, target) in enumerate(train_loader):
data, target = data.to(device), target.to(device) # 将数据和目标迁移到 CUDA
optimizer.zero_grad()
output = model(data)
target_last = target - data[:, 24:36, :, :, :]
# loss = criterion(output, target_last) # 使用 L1 损失
loss = criterion(output[:,:12,:,:71,:], target_last[:,:12,:,:71,:]) # 使用 L1 损失
print(loss)
all_loss+=loss
loss.backward()
optimizer.step()
print(f'Epoch {epoch + 1}/{max_epochs}, Train Loss: {all_loss.item():.4f}')
# 验证阶段
model.eval()
val_loss = 0.0
with torch.no_grad():
for data, target in val_loader:
data, target = data.to(device), target.to(device) # 将数据和目标迁移到 CUDA
output = model(data)
target_last = target - data[:, 24:36, :, :, :]
# loss = criterion(output, target_last) # 使用 L1 损失
loss = criterion(output[:, :12, :, :71, :], target_last[:, :12, :, :71, :]) # 使用 L1 损失
val_loss += loss.item()
val_loss /= len(val_loader)
print(f'Epoch {epoch + 1}/{max_epochs}, Val Loss: {val_loss:.4f}')
# 保存最佳模型
if val_loss < best_val_loss:
best_val_loss = val_loss
torch.save(model.state_dict(), 'best_model.pth')
print('Best model saved!')
print('Training completed.') |