{
"cells": [
{
"cell_type": "code",
"execution_count": 1,
"id": "3a2c33ed-8f78-4ce4-b5cd-7b7ffc5c8273",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Using device: cuda\n"
]
}
],
"source": [
"import os\n",
"import numpy as np\n",
"import pandas as pd\n",
"import torch\n",
"import torch.nn as nn\n",
"import torch.optim as optim\n",
"from torch.utils.data import DataLoader, Dataset\n",
"from PIL import Image\n",
"import matplotlib.pyplot as plt\n",
"os.environ[\"CUDA_VISIBLE_DEVICE\"] = \"0\" \n",
"\n",
"\n",
"# 设置CUDA设备\n",
"device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
"print(f\"Using device: {device}\")"
]
},
{
"cell_type": "code",
"execution_count": 2,
"id": "85459fd1-6835-41cd-b645-553611c358e8",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Maximum pixel value in the dataset: 107.49169921875\n"
]
}
],
"source": [
"max_pixel_value = 107.49169921875\n",
"\n",
"print(f\"Maximum pixel value in the dataset: {max_pixel_value}\")"
]
},
{
"cell_type": "code",
"execution_count": 3,
"id": "3fc0918e-103c-40a3-93bc-6171e934a7e4",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"checkpoint before Generator is OK\n"
]
}
],
"source": [
"class NO2Dataset(Dataset):\n",
" \n",
" def __init__(self, image_dir, mask_dir):\n",
" \n",
" self.image_dir = image_dir\n",
" self.mask_dir = mask_dir\n",
" self.image_filenames = [f for f in os.listdir(image_dir) if f.endswith('.npy')] # 仅加载 .npy 文件\n",
" self.mask_filenames = [f for f in os.listdir(mask_dir) if f.endswith('.jpg')] # 仅加载 .jpg 文件\n",
" \n",
" def __len__(self):\n",
" \n",
" return len(self.image_filenames)\n",
" \n",
" def __getitem__(self, idx):\n",
" \n",
" image_path = os.path.join(self.image_dir, self.image_filenames[idx])\n",
" mask_idx = idx % len(self.mask_filenames)\n",
" mask_path = os.path.join(self.mask_dir, self.mask_filenames[mask_idx])\n",
"\n",
" # 加载图像数据 (.npy 文件)\n",
" image = np.load(image_path).astype(np.float32) / max_pixel_value # 形状为 (96, 96, 8)\n",
"\n",
" # 加载掩码数据 (.jpg 文件)\n",
" mask = np.array(Image.open(mask_path).convert('L')).astype(np.float32)\n",
"\n",
" # 将掩码数据中非0值设为1，0值保持不变\n",
" mask = np.where(mask != 0, 1.0, 0.0)\n",
"\n",
" # 保持掩码数据形状为 (96, 96, 1)\n",
" mask = mask[:, :, np.newaxis] # 将形状调整为 (96, 96, 1)\n",
"\n",
" # 应用掩码\n",
" masked_image = image.copy()\n",
" masked_image[:, :, 0] = image[:, :, 0] * mask.squeeze() # 遮盖NO2数据\n",
"\n",
" # cGAN的输入和目标\n",
" X = np.concatenate([masked_image[:, :, :1], image[:, :, 1:]], axis=-1) # 形状为 (96, 96, 8)\n",
" y = image[:, :, 0:1] # 目标输出为NO2数据，形状为 (96, 96, 1)\n",
"\n",
" # 转换形状为 (channels, height, width)\n",
" X = np.transpose(X, (2, 0, 1)) # 转换为 (8, 96, 96)\n",
" y = np.transpose(y, (2, 0, 1)) # 转换为 (1, 96, 96)\n",
" mask = np.transpose(mask, (2, 0, 1)) # 转换为 (1, 96, 96)\n",
"\n",
" return torch.tensor(X, dtype=torch.float32), torch.tensor(y, dtype=torch.float32), torch.tensor(mask, dtype=torch.float32)\n",
"\n",
"# 实例化数据集和数据加载器\n",
"train_dir = './out_mat/96/train/'\n",
"valid_dir = './out_mat/96/valid/'\n",
"test_dir = './out_mat/96/test/'\n",
"mask_dir = './out_mat/96/mask/20/'\n",
"\n",
"print(f\"checkpoint before Generator is OK\")\n",
"\n",
"dataset = NO2Dataset(train_dir, mask_dir)\n",
"train_loader = DataLoader(dataset, batch_size=64, shuffle=True, num_workers=8)\n",
"\n",
"validset = NO2Dataset(valid_dir, mask_dir)\n",
"val_loader = DataLoader(validset, batch_size=64, shuffle=False, num_workers=8)\n",
"\n",
"testset = NO2Dataset(test_dir, mask_dir)\n",
"test_loader = DataLoader(testset, batch_size=64, shuffle=False, num_workers=8)"
]
},
{
"cell_type": "code",
"execution_count": 4,
"id": "a60b7019-f231-4ccb-9195-c459f3a1521d",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Generator is on: cuda:0\n",
"Discriminator is on: cuda:0\n"
]
}
],
"source": [
"# 生成器模型\n",
"class Generator(nn.Module):\n",
" \n",
" def __init__(self):\n",
" super(Generator, self).__init__()\n",
" self.encoder = nn.Sequential(\n",
" nn.Conv2d(8, 64, kernel_size=4, stride=2, padding=1),\n",
" nn.LeakyReLU(0.2, inplace=True),\n",
" nn.BatchNorm2d(64),\n",
" nn.Conv2d(64, 128, kernel_size=4, stride=2, padding=1),\n",
" nn.LeakyReLU(0.2, inplace=True),\n",
" nn.BatchNorm2d(128),\n",
" )\n",
" self.decoder = nn.Sequential(\n",
" nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1),\n",
" nn.ReLU(inplace=True),\n",
" nn.BatchNorm2d(64),\n",
" nn.ConvTranspose2d(64, 1, kernel_size=4, stride=2, padding=1),\n",
" nn.Tanh(),\n",
" )\n",
"\n",
" def forward(self, x):\n",
" x_encoded = self.encoder(x)\n",
" x_decoded = self.decoder(x_encoded)\n",
"\n",
"# x_decoded = (x_decoded + 1) / 2\n",
"\n",
"# x_output = (1 - mask) * x_decoded + mask * x[:, :1, :, :]\n",
" return x_output\n",
"\n",
"# 判别器模型\n",
"class Discriminator(nn.Module):\n",
" \n",
" def __init__(self):\n",
" super(Discriminator, self).__init__()\n",
" self.model = nn.Sequential(\n",
" nn.Conv2d(1, 64, kernel_size=4, stride=2, padding=1),\n",
" nn.LeakyReLU(0.2, inplace=True),\n",
" nn.BatchNorm2d(64),\n",
" nn.Conv2d(64, 128, kernel_size=4, stride=2, padding=1),\n",
" nn.LeakyReLU(0.2, inplace=True),\n",
" nn.BatchNorm2d(128),\n",
" nn.Conv2d(128, 1, kernel_size=4, stride=2, padding=1),\n",
" nn.Sigmoid(),\n",
" )\n",
"\n",
" def forward(self, x):\n",
" return self.model(x)\n",
"\n",
"# 将模型加载到GPU\n",
"generator = Generator().to(device)\n",
"discriminator = Discriminator().to(device)\n",
"\n",
"# 定义优化器和损失函数\n",
"optimizer_G = optim.Adam(generator.parameters(), lr=0.0002, betas=(0.5, 0.999))\n",
"optimizer_D = optim.Adam(discriminator.parameters(), lr=0.0002, betas=(0.5, 0.999))\n",
"adversarial_loss = nn.BCELoss().to(device)\n",
"pixelwise_loss = nn.MSELoss().to(device)\n",
"\n",
"# 确认模型是否在GPU上\n",
"print(f\"Generator is on: {next(generator.parameters()).device}\")\n",
"print(f\"Discriminator is on: {next(discriminator.parameters()).device}\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "645dd325-fc70-4234-8279-bc8cbc4c5dde",
"metadata": {},
"outputs": [],
"source": [
"# 开始训练\n",
"epochs = 150\n",
"for epoch in range(epochs):\n",
" for i, (X, y, mask) in enumerate(train_loader):\n",
" # 将数据移到 GPU 上\n",
" X, y, mask = X.to(device), y.to(device), mask.to(device)\n",
" # print(f\"X is on: {X.device}, y is on: {y.device}, mask is on: {mask.device}, i = {i}\") #checkpoint\n",
" \n",
" valid = torch.ones((X.size(0), 1, 12, 12)).to(device)\n",
" fake = torch.zeros((X.size(0), 1, 12, 12)).to(device)\n",
"\n",
" # 生成器生成图像\n",
" optimizer_G.zero_grad()\n",
" generated_images = generator(X, mask)\n",
" g_loss = adversarial_loss(discriminator(torch.cat((generated_images, X), dim=1)), valid) + 100 * pixelwise_loss(\n",
" generated_images, y)\n",
" g_loss.backward()\n",
" optimizer_G.step()\n",
"\n",
" # 判别器训练\n",
" optimizer_D.zero_grad()\n",
" real_loss = adversarial_loss(discriminator(torch.cat((y, X), dim=1)), valid)\n",
" fake_loss = adversarial_loss(discriminator(torch.cat((generated_images.detach(), X), dim=1)), fake)\n",
" d_loss = 0.5 * (real_loss + fake_loss)\n",
" d_loss.backward()\n",
" optimizer_D.step()\n",
"\n",
" print(f\"Epoch [{epoch}/{epochs}] [D loss: {d_loss.item()}] [G loss: {g_loss.item()}]\")\n",
"\n",
"# 保存训练好的模型\n",
"torch.save(generator.state_dict(), './models/GAN/generator.pth')\n",
"torch.save(discriminator.state_dict(), './models/GAN/discriminator.pth')"
]
},
{
"cell_type": "code",
"execution_count": 9,
"id": "2d889a53-5415-4895-99ff-fc63745884a5",
"metadata": {},
"outputs": [],
"source": [
"from sklearn.metrics import mean_squared_error, r2_score, mean_absolute_percentage_error, mean_absolute_error"
]
},
{
"cell_type": "code",
"execution_count": 10,
"id": "37f2df19-492c-4231-a388-13182ce515db",
"metadata": {},
"outputs": [],
"source": [
"def cal_ioa(y_true, y_pred):\n",
" # 计算平均值\n",
" mean_observed = np.mean(y_true)\n",
" mean_predicted = np.mean(y_pred)\n",
"\n",
" # 计算IoA\n",
" numerator = np.sum((y_true - y_pred) ** 2)\n",
" denominator = np.sum((np.abs(y_true - mean_observed) + np.abs(y_pred - mean_predicted)) ** 2)\n",
" IoA = 1 - (numerator / denominator)\n",
"\n",
" return IoA"
]
},
{
"cell_type": "code",
"execution_count": 11,
"id": "0b93a2e7-c4fb-4611-9967-4e33f9982ad5",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"eva_list = list()\n",
"device = 'cpu'\n",
"generator = generator.to(device)\n",
"with torch.no_grad():\n",
" for batch_idx, (X, y, mask) in enumerate(test_loader):\n",
" X, y, mask = X.to(device), y.to(device), mask.to(device)\n",
" mask_rev = (torch.squeeze(mask, dim=1)==0) * 1 # mask取反获得修复区域\n",
" reconstructed = generator(X, mask)\n",
" rev_data = torch.squeeze(y * max_pixel_value, dim=1)\n",
" rev_recon = torch.squeeze(reconstructed * max_pixel_value, dim=1)\n",
" # todo: 这里需要只评估修补出来的模块\n",
" data_label = rev_data * mask_rev\n",
" data_label = data_label[mask_rev==1]\n",
" recon_no2 = rev_recon * mask_rev\n",
" recon_no2 = recon_no2[mask_rev==1]\n",
" y_true = rev_data.flatten()\n",
" y_pred = rev_recon.flatten()\n",
" mae = mean_absolute_error(y_true, y_pred)\n",
" rmse = np.sqrt(mean_squared_error(y_true, y_pred))\n",
" mape = mean_absolute_percentage_error(y_true, y_pred)\n",
" r2 = r2_score(y_true, y_pred)\n",
" ioa = cal_ioa(data_label.detach().numpy(), recon_no2.detach().numpy())\n",
" eva_list.append([mae, rmse, mape, r2, ioa])"
]
},
{
"cell_type": "code",
"execution_count": 22,
"id": "e7b3323f-7116-4d4e-8483-2fd605e2fb57",
"metadata": {},
"outputs": [
{
"data": {
"text/html": [
"\n",
"\n",
"
\n",
" \n",
" \n",
" | \n",
" mae | \n",
" rmse | \n",
" mape | \n",
" r2 | \n",
" ioa | \n",
"
\n",
" \n",
" \n",
" \n",
" | count | \n",
" 75.000000 | \n",
" 75.000000 | \n",
" 75.000000 | \n",
" 75.000000 | \n",
" 75.000000 | \n",
"
\n",
" \n",
" | mean | \n",
" 0.399366 | \n",
" 1.246761 | \n",
" 0.047188 | \n",
" 0.963991 | \n",
" 0.939587 | \n",
"
\n",
" \n",
" | std | \n",
" 0.071295 | \n",
" 0.220616 | \n",
" 0.005035 | \n",
" 0.018081 | \n",
" 0.026807 | \n",
"
\n",
" \n",
" | min | \n",
" 0.348072 | \n",
" 1.073966 | \n",
" 0.040716 | \n",
" 0.813181 | \n",
" 0.719442 | \n",
"
\n",
" \n",
" | 25% | \n",
" 0.372762 | \n",
" 1.154684 | \n",
" 0.043751 | \n",
" 0.963354 | \n",
" 0.938713 | \n",
"
\n",
" \n",
" | 50% | \n",
" 0.388768 | \n",
" 1.207949 | \n",
" 0.045860 | \n",
" 0.966356 | \n",
" 0.943430 | \n",
"
\n",
" \n",
" | 75% | \n",
" 0.402351 | \n",
" 1.274836 | \n",
" 0.050051 | \n",
" 0.968919 | \n",
" 0.947026 | \n",
"
\n",
" \n",
" | max | \n",
" 0.959251 | \n",
" 2.966476 | \n",
" 0.066256 | \n",
" 0.972840 | \n",
" 0.956998 | \n",
"
\n",
" \n",
"
\n",
"
\n",
"\n",
"
\n",
" \n",
" \n",
" | \n",
" mae | \n",
" rmse | \n",
" mape | \n",
" r2 | \n",
" ioa | \n",
" r | \n",
"
\n",
" \n",
" \n",
" \n",
" | count | \n",
" 4739.000000 | \n",
" 4739.000000 | \n",
" 4739.000000 | \n",
" 4739.000000 | \n",
" 4739.000000 | \n",
" 4739.000000 | \n",
"
\n",
" \n",
" | mean | \n",
" 1.977856 | \n",
" 2.505394 | \n",
" 0.236766 | \n",
" 0.444234 | \n",
" 0.826185 | \n",
" 0.795505 | \n",
"
\n",
" \n",
" | std | \n",
" 0.966137 | \n",
" 1.156947 | \n",
" 0.075139 | \n",
" 0.309037 | \n",
" 0.112186 | \n",
" 0.109227 | \n",
"
\n",
" \n",
" | min | \n",
" 0.588599 | \n",
" 0.782554 | \n",
" 0.106112 | \n",
" -5.779783 | \n",
" -2.754070 | \n",
" 0.284676 | \n",
"
\n",
" \n",
" | 25% | \n",
" 1.195541 | \n",
" 1.551567 | \n",
" 0.187231 | \n",
" 0.300401 | \n",
" 0.781712 | \n",
" 0.735376 | \n",
"
\n",
" \n",
" | 50% | \n",
" 1.606092 | \n",
" 2.094027 | \n",
" 0.220013 | \n",
" 0.506733 | \n",
" 0.849549 | \n",
" 0.822590 | \n",
"
\n",
" \n",
" | 75% | \n",
" 2.658243 | \n",
" 3.338708 | \n",
" 0.266574 | \n",
" 0.658528 | \n",
" 0.899010 | \n",
" 0.876813 | \n",
"
\n",
" \n",
" | max | \n",
" 9.428754 | \n",
" 9.982598 | \n",
" 0.903847 | \n",
" 0.889351 | \n",
" 0.969285 | \n",
" 0.960868 | \n",
"
\n",
" \n",
"
\n",
"