MAE_ATMO/torch_MAE_1d_baseline.ipynb

In [30]:

import os
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset, random_split
from PIL import Image
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import cv2

In [4]:

np.random.seed(0)
torch.random.manual_seed(0)

Out[4]:

<torch._C.Generator at 0x7f6d6be638f0>

In [5]:

# 定义函数来找到最大值
def find_max_pixel_value(image_dir):
    max_pixel_value = 0.0
    for filename in os.listdir(image_dir):
        if filename.endswith('.npy'):
            image_path = os.path.join(image_dir, filename)
            image = np.load(image_path).astype(np.float32)
            max_pixel_value = max(max_pixel_value, image[:, :, 0].max())
    return max_pixel_value

# 计算图像数据中的最大像素值
image_dir = './out_mat/96/train/' 
max_pixel_value = find_max_pixel_value(image_dir)

print(f"Maximum pixel value in the dataset: {max_pixel_value}")

Maximum pixel value in the dataset: 107.49169921875

In [6]:

class NO2Dataset(Dataset):
    
    def __init__(self, image_dir, mask_dir):
        
        self.image_dir = image_dir
        self.mask_dir = mask_dir
        self.image_filenames = [f for f in os.listdir(image_dir) if f.endswith('.npy')]  # 仅加载 .npy 文件
        self.mask_filenames = [f for f in os.listdir(mask_dir) if f.endswith('.jpg')]  # 仅加载 .jpg 文件
        
    def __len__(self):
        
        return len(self.image_filenames)
    
    def __getitem__(self, idx):
        
        image_path = os.path.join(self.image_dir, self.image_filenames[idx])
        mask_idx = np.random.choice(self.mask_filenames)
        mask_path = os.path.join(self.mask_dir, mask_idx)

        # 加载图像数据 (.npy 文件)
        image = np.load(image_path).astype(np.float32)[:,:,:1] / max_pixel_value  # 形状为 (96, 96, 1)

        # 加载掩码数据 (.jpg 文件)
        mask = np.array(Image.open(mask_path).convert('L')).astype(np.float32)

        # 将掩码数据中非0值设为1，0值保持不变
        mask = np.where(mask != 0, 1.0, 0.0)

        # 保持掩码数据形状为 (96, 96, 1)
        mask = mask[:, :, np.newaxis]  # 将形状调整为 (96, 96, 1)

        # 应用掩码
        masked_image = image.copy()
        masked_image[:, :, 0] = image[:, :, 0] * mask.squeeze()  # 遮盖NO2数据

        # cGAN的输入和目标
        X = masked_image[:, :, :1]  # 形状为 (96, 96, 8)
        y = image[:, :, 0:1]  # 目标输出为NO2数据，形状为 (96, 96, 1)

        # 转换形状为 (channels, height, width)
        X = np.transpose(X, (2, 0, 1))  # 转换为 (1, 96, 96)
        y = np.transpose(y, (2, 0, 1))  # 转换为 (1, 96, 96)
        mask = np.transpose(mask, (2, 0, 1))  # 转换为 (1, 96, 96)

        return torch.tensor(X, dtype=torch.float32), torch.tensor(y, dtype=torch.float32), torch.tensor(mask, dtype=torch.float32)

# 实例化数据集和数据加载器
image_dir = './out_mat/96/train/'
mask_dir = './out_mat/96/mask/20/'

print(f"checkpoint before Generator is OK")

checkpoint before Generator is OK

In [7]:

dataset = NO2Dataset(image_dir, mask_dir)
dataloader = DataLoader(dataset, batch_size=64, shuffle=True, num_workers=8)
val_set = NO2Dataset('./out_mat/96/valid/', mask_dir)
val_loader = DataLoader(val_set, batch_size=64, shuffle=False, num_workers=4)
test_set = NO2Dataset('./out_mat/96/test/', mask_dir)
test_loader = DataLoader(test_set, batch_size=64, shuffle=False, num_workers=4)

In [8]:

# 可视化特定特征的函数
def visualize_feature(input_feature,masked_feature, output_feature, title):
    plt.figure(figsize=(12, 6))
    plt.subplot(1, 3, 1)
    plt.imshow(input_feature[0].cpu().numpy(), cmap='RdYlGn_r')
    plt.title(title + " Input")
    plt.subplot(1, 3, 2)
    plt.imshow(masked_feature[0].cpu().numpy(), cmap='RdYlGn_r')
    plt.title(title + " Masked")
    plt.subplot(1, 3, 3)
    plt.imshow(output_feature[0].detach().cpu().numpy(), cmap='RdYlGn_r')
    plt.title(title + " Recovery")
    plt.show()

In [9]:

# 设置随机种子以确保结果的可重复性
torch.manual_seed(0)
np.random.seed(0)

# 数据准备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(device)

cuda

In [10]:

# 定义Masked Autoencoder模型
class MaskedAutoencoderBase(nn.Module):
    def __init__(self):
        super(MaskedAutoencoderBase, self).__init__()
        self.encoder = nn.Sequential(
            nn.Conv2d(1, 32, kernel_size=3, stride=2, padding=1),
            nn.ReLU(),
            nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1),
            nn.ReLU(),
            nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1),
            nn.ReLU(),
        )
        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(128, 32, kernel_size=3, stride=2, padding=1, output_padding=1),
            nn.ReLU(),
            nn.ConvTranspose2d(32, 16, kernel_size=3, stride=2, padding=1, output_padding=1),
            nn.ReLU(),
            nn.ConvTranspose2d(16, 1, kernel_size=3, stride=2, padding=1, output_padding=1),
            nn.Sigmoid()  # 使用Sigmoid是因为输入数据是0-1之间的
        )

    def forward(self, x):
        encoded = self.encoder(x)
        decoded = self.decoder(encoded)
        return decoded

In [11]:

def masked_mse_loss(preds, target, mask):
    loss = (preds - target) ** 2
    loss = loss.mean(dim=-1)  # 对每个像素点求平均
    loss = (loss * mask).sum() / mask.sum()  # 只计算被mask的像素点的损失
    return loss

In [12]:

# 实例化模型、损失函数和优化器
model = MaskedAutoencoderBase()
criterion = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)

In [18]:

# 训练函数
def train_epoch(model, device, data_loader, criterion, optimizer):
    model.train()
    running_loss = 0.0
    for batch_idx, (X, y, mask) in enumerate(data_loader):
        X, y, mask = X.to(device), y.to(device), mask.to(device)
        optimizer.zero_grad()
        reconstructed = model(X)
        loss = masked_mse_loss(reconstructed, y, mask)
        # loss = criterion(reconstructed, y)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
    return running_loss / (batch_idx + 1)

In [19]:

# 评估函数
def evaluate(model, device, data_loader, criterion):
    model.eval()
    running_loss = 0.0
    with torch.no_grad():
        for batch_idx, (X, y, mask) in enumerate(data_loader):
            X, y, mask = X.to(device), y.to(device), mask.to(device)
            reconstructed = model(X)
            if batch_idx == 8:
                rand_ind = np.random.randint(0, len(y))
                # visualize_feature(y[rand_ind], X[rand_ind], reconstructed[rand_ind], title='NO_2')
            loss = masked_mse_loss(reconstructed, y, mask)
            # loss = criterion(reconstructed, y)
            running_loss += loss.item()
    return running_loss / (batch_idx + 1)

In [20]:

model = model.to(device)

num_epochs = 100
train_losses = list()
val_losses = list()
for epoch in range(num_epochs):
    train_loss = train_epoch(model, device, dataloader, criterion, optimizer)
    train_losses.append(train_loss)
    val_loss = evaluate(model, device, val_loader, criterion)
    val_losses.append(val_loss)
    print(f'Epoch {epoch+1}, Train Loss: {train_loss}, Val Loss: {val_loss}')

# 测试模型
test_loss = evaluate(model, device, test_loader, criterion)
print(f'Test Loss: {test_loss}')

Epoch 1, Train Loss: 2.4377448216936233, Val Loss: 0.1723788405087457
Epoch 2, Train Loss: 0.09637997932197374, Val Loss: 0.07621741728551353
Epoch 3, Train Loss: 0.06397618102934657, Val Loss: 0.06195200451496822
Epoch 4, Train Loss: 0.052692288621974906, Val Loss: 0.052603201690449644
Epoch 5, Train Loss: 0.045533701719529036, Val Loss: 0.0462518873721806
Epoch 6, Train Loss: 0.040426999678095564, Val Loss: 0.04118765834996949
Epoch 7, Train Loss: 0.03643315702979787, Val Loss: 0.0370612932619319
Epoch 8, Train Loss: 0.03297993362074691, Val Loss: 0.0338741072189452
Epoch 9, Train Loss: 0.030229569176595177, Val Loss: 0.03180063916231269
Epoch 10, Train Loss: 0.028299911767600827, Val Loss: 0.03058352780097456
Epoch 11, Train Loss: 0.026935207724357337, Val Loss: 0.029766072282817826
Epoch 12, Train Loss: 0.026076676769618782, Val Loss: 0.028107319638800265
Epoch 13, Train Loss: 0.02534967841821139, Val Loss: 0.027272115475428637
Epoch 14, Train Loss: 0.024701394381349166, Val Loss: 0.02684043228292643
Epoch 15, Train Loss: 0.0240272392550011, Val Loss: 0.02594853615138068
Epoch 16, Train Loss: 0.0233813104438083, Val Loss: 0.025640942656726978
Epoch 17, Train Loss: 0.02310016915273438, Val Loss: 0.02571806650775582
Epoch 18, Train Loss: 0.022718923658792054, Val Loss: 0.024644668200122778
Epoch 19, Train Loss: 0.022323213453052576, Val Loss: 0.024273945435659208
Epoch 20, Train Loss: 0.02199719715685223, Val Loss: 0.02410240029332353
Epoch 21, Train Loss: 0.021530815467024535, Val Loss: 0.02380427871066243
Epoch 22, Train Loss: 0.021460241776262743, Val Loss: 0.0232450627346537
Epoch 23, Train Loss: 0.02090326771050977, Val Loss: 0.022885078564286232
Epoch 24, Train Loss: 0.020652044475363774, Val Loss: 0.022562191390724323
Epoch 25, Train Loss: 0.02051923798985387, Val Loss: 0.022203324724044373
Epoch 26, Train Loss: 0.020149177833767743, Val Loss: 0.022744494337421744
Epoch 27, Train Loss: 0.020068248640300268, Val Loss: 0.022425833088693333
Epoch 28, Train Loss: 0.019720358143529397, Val Loss: 0.02253118777341807
Epoch 29, Train Loss: 0.01939903690288084, Val Loss: 0.021765351378873213
Epoch 30, Train Loss: 0.01943497322989922, Val Loss: 0.021345259649540062
Epoch 31, Train Loss: 0.019241397384928458, Val Loss: 0.02124041018646155
Epoch 32, Train Loss: 0.01906546402464144, Val Loss: 0.021633521083797982
Epoch 33, Train Loss: 0.01884070100512302, Val Loss: 0.021043253979131357
Epoch 34, Train Loss: 0.01874133140855785, Val Loss: 0.02059999839472237
Epoch 35, Train Loss: 0.01853996916544851, Val Loss: 0.021178998303279947
Epoch 36, Train Loss: 0.018260161060412106, Val Loss: 0.020367807639178944
Epoch 37, Train Loss: 0.01830708983233956, Val Loss: 0.020017842692670536
Epoch 38, Train Loss: 0.018042967790675362, Val Loss: 0.020187884722071798
Epoch 39, Train Loss: 0.017922898732197056, Val Loss: 0.019615614786744118
Epoch 40, Train Loss: 0.017794321282236486, Val Loss: 0.019430582606191956
Epoch 41, Train Loss: 0.017688655022656517, Val Loss: 0.019477688401603875
Epoch 42, Train Loss: 0.017460078103512383, Val Loss: 0.018902005530448993
Epoch 43, Train Loss: 0.01727662416638441, Val Loss: 0.018832763184362382
Epoch 44, Train Loss: 0.017280888195599666, Val Loss: 0.019056980081124983
Epoch 45, Train Loss: 0.017114856775012312, Val Loss: 0.018604515495696174
Epoch 46, Train Loss: 0.016909640970858234, Val Loss: 0.018437264904157438
Epoch 47, Train Loss: 0.016691252999185946, Val Loss: 0.01889144025965413
Epoch 48, Train Loss: 0.016869753608079047, Val Loss: 0.018732781104965887
Epoch 49, Train Loss: 0.01653263871179243, Val Loss: 0.01850963812043418
Epoch 50, Train Loss: 0.01653244017520875, Val Loss: 0.0178856217344083
Epoch 51, Train Loss: 0.016499577624874823, Val Loss: 0.01781756227919415
Epoch 52, Train Loss: 0.016335643743249504, Val Loss: 0.01821571894323648
Epoch 53, Train Loss: 0.016375035212406415, Val Loss: 0.017511379168327176
Epoch 54, Train Loss: 0.016288986672428948, Val Loss: 0.017456448650849398
Epoch 55, Train Loss: 0.01623404509517137, Val Loss: 0.017827068525018978
Epoch 56, Train Loss: 0.016188283936615196, Val Loss: 0.017475027326883663
Epoch 57, Train Loss: 0.01605349867359588, Val Loss: 0.017256822728955033
Epoch 58, Train Loss: 0.015958637990610022, Val Loss: 0.017457437256712522
Epoch 59, Train Loss: 0.016034694237001774, Val Loss: 0.017437012713235705
Epoch 60, Train Loss: 0.0158486066956483, Val Loss: 0.017560158175096582
Epoch 61, Train Loss: 0.015632042563275286, Val Loss: 0.01692103194211846
Epoch 62, Train Loss: 0.015540152108608677, Val Loss: 0.01698271286632143
Epoch 63, Train Loss: 0.01545496231043025, Val Loss: 0.01699626362368242
Epoch 64, Train Loss: 0.015430795162488398, Val Loss: 0.01687317063559347
Epoch 65, Train Loss: 0.015489797350732191, Val Loss: 0.017046043955123248
Epoch 66, Train Loss: 0.015236956011682179, Val Loss: 0.0172197060214717
Epoch 67, Train Loss: 0.015348140916755895, Val Loss: 0.016508253249548265
Epoch 68, Train Loss: 0.015228347097519055, Val Loss: 0.016413471842212462
Epoch 69, Train Loss: 0.01516882229025997, Val Loss: 0.01686259738600521
Epoch 70, Train Loss: 0.015173258315593574, Val Loss: 0.01757873013726811
Epoch 71, Train Loss: 0.015156847716678986, Val Loss: 0.016662339123883353
Epoch 72, Train Loss: 0.015105586064507088, Val Loss: 0.016890839868183457
Epoch 73, Train Loss: 0.014925161955887051, Val Loss: 0.015931842709655194
Epoch 74, Train Loss: 0.014886363126497947, Val Loss: 0.016006485308840204
Epoch 75, Train Loss: 0.015015289608531735, Val Loss: 0.015968994154080526
Epoch 76, Train Loss: 0.014806462892968403, Val Loss: 0.015919692327838336
Epoch 77, Train Loss: 0.014728168116962653, Val Loss: 0.015852669684855797
Epoch 78, Train Loss: 0.014845167781319914, Val Loss: 0.016079049404543726
Epoch 79, Train Loss: 0.014719554133998435, Val Loss: 0.015957326447563387
Epoch 80, Train Loss: 0.014635249268281404, Val Loss: 0.015849308388780303
Epoch 81, Train Loss: 0.014474964379800849, Val Loss: 0.015526832887597049
Epoch 82, Train Loss: 0.014369143295641007, Val Loss: 0.015485089967277512
Epoch 83, Train Loss: 0.014446225396076743, Val Loss: 0.015848276135859204
Epoch 84, Train Loss: 0.014476079110537419, Val Loss: 0.015343323600158762
Epoch 85, Train Loss: 0.014672522836378174, Val Loss: 0.015515949938501885
Epoch 86, Train Loss: 0.014440825409545568, Val Loss: 0.015224411166203556
Epoch 87, Train Loss: 0.014462759978980111, Val Loss: 0.015663697370397512
Epoch 88, Train Loss: 0.01440465696262971, Val Loss: 0.015856551353944773
Epoch 89, Train Loss: 0.014255739579146559, Val Loss: 0.015246662380757616
Epoch 90, Train Loss: 0.014205876624202756, Val Loss: 0.015011716536732752
Epoch 91, Train Loss: 0.014259663818216924, Val Loss: 0.015085076996639593
Epoch 92, Train Loss: 0.014251617286978156, Val Loss: 0.015133185506756627
Epoch 93, Train Loss: 0.014119144052302723, Val Loss: 0.015415464166496227
Epoch 94, Train Loss: 0.014192042264053554, Val Loss: 0.015254960033986995
Epoch 95, Train Loss: 0.014140318196855094, Val Loss: 0.017451592276234235
Epoch 96, Train Loss: 0.014092271857890502, Val Loss: 0.015359595265072672
Epoch 97, Train Loss: 0.01409529693843574, Val Loss: 0.015055305060388437
Epoch 98, Train Loss: 0.014136464546688578, Val Loss: 0.015083992547953307
Epoch 99, Train Loss: 0.013914715411792103, Val Loss: 0.014718598477653604
Epoch 100, Train Loss: 0.013870610982518305, Val Loss: 0.01483591334588492
Test Loss: 0.010182651478874807

In [24]:

tr_ind = list(range(len(train_losses)))
val_ind = list(range(len(val_losses)))
plt.plot(train_losses, label='train_loss')
plt.plot(val_losses, label='val_loss')
plt.legend(loc='best')

Out[24]:

<matplotlib.legend.Legend at 0x7f6b7999d190>

In [25]:

from sklearn.metrics import mean_squared_error, r2_score, mean_absolute_percentage_error, mean_absolute_error

In [26]:

eva_list = list()
device = 'cpu'
model = model.to(device)
with torch.no_grad():
    for batch_idx, (X, y, mask) in enumerate(test_loader):
        X, y, mask = X.to(device), y.to(device), mask.to(device)
        mask_rev = (torch.squeeze(mask, dim=1)==0) * 1 # mask取反获得修复区域
        reconstructed = model(X)
        # tr_maxs = np.transpose(maxs, (2, 0, 1))
        # tr_mins = np.transpose(mins, (2, 0, 1))
        rev_data = y * max_pixel_value
        rev_recon = reconstructed * max_pixel_value
        # todo: 这里需要只评估修补出来的模块
        data_label = torch.squeeze(rev_data, dim=1) * mask_rev
        data_label = data_label[mask_rev==1]
        recon_no2 = torch.squeeze(rev_recon, dim=1) * mask_rev
        recon_no2 = recon_no2[mask_rev==1]
        mae = mean_absolute_error(data_label, recon_no2)
        rmse = np.sqrt(mean_squared_error(data_label, recon_no2))
        mape = mean_absolute_percentage_error(data_label, recon_no2)
        r2 = r2_score(data_label, recon_no2)
        eva_list.append([mae, rmse, mape, r2])

In [28]:

pd.DataFrame(eva_list, columns=['mae', 'rmse', 'mape', 'r2']).describe()

Out[28]:

	mae	rmse	mape	r2
count	75.000000	75.000000	75.000000	75.000000
mean	1.669181	2.722375	0.228690	0.825025
std	0.101549	0.229373	0.027824	0.023960
min	1.456919	2.206495	0.147438	0.751642
25%	1.600787	2.569844	0.210564	0.815437
50%	1.663539	2.723380	0.228493	0.826285
75%	1.726697	2.848122	0.248380	0.837574
max	1.998206	3.443690	0.287797	0.881901

In [31]:

def cal_ioa(y_true, y_pred):
    # 计算平均值
    mean_observed = np.mean(y_true)
    mean_predicted = np.mean(y_pred)

    # 计算IoA
    numerator = np.sum((y_true - y_pred) ** 2)
    denominator = np.sum((np.abs(y_true - mean_observed) + np.abs(y_pred - mean_predicted)) ** 2)
    IoA = 1 - (numerator / denominator)

    return IoA

In [33]:

eva_list_frame = list()
device = 'cpu'
model = model.to(device)
with torch.no_grad():
    for batch_idx, (X, y, mask) in enumerate(test_loader):
        X, y, mask = X.to(device), y.to(device), mask.to(device)
        mask_rev = (torch.squeeze(mask, dim=1)==0) * 1 # mask取反获得修复区域
        reconstructed = model(X)
        rev_data = y * max_pixel_value
        rev_recon = reconstructed * max_pixel_value
        # todo: 这里需要只评估修补出来的模块
        for i, sample in enumerate(rev_data):
            used_mask = mask_rev[i]
            data_label = sample[0] * used_mask
            recon_no2 = rev_recon[i][0] * used_mask
            data_label = data_label[used_mask==1]
            recon_no2 = recon_no2[used_mask==1]
            mae = mean_absolute_error(data_label, recon_no2)
            rmse = np.sqrt(mean_squared_error(data_label, recon_no2))
            mape = mean_absolute_percentage_error(data_label, recon_no2)
            r2 = r2_score(data_label, recon_no2)
            ioa = cal_ioa(data_label.detach().numpy(), recon_no2.detach().numpy())
            r = np.corrcoef(data_label, recon_no2)[0, 1]
            eva_list_frame.append([mae, rmse, mape, r2, ioa, r])

In [34]:

pd.DataFrame(eva_list_frame, columns=['mae', 'rmse', 'mape', 'r2', 'ioa', 'r']).describe()

Out[34]:

	mae	rmse	mape	r2	ioa	r
count	4739.000000	4739.000000	4739.000000	4739.000000	4739.000000	4739.000000
mean	1.657513	2.352886	0.232260	0.406715	0.823672	0.747644
std	0.897659	1.318793	0.234080	0.877368	0.184708	0.191853
min	0.546354	0.695038	0.066870	-30.315991	-1.254103	-0.392216
25%	1.042898	1.472388	0.137216	0.313886	0.782728	0.671206
50%	1.465436	2.072718	0.174430	0.610684	0.879875	0.805015
75%	1.976618	2.785021	0.234618	0.757136	0.929170	0.879403
max	9.007959	12.398485	3.290891	0.973600	0.993247	0.987535

In [29]:

pd.DataFrame(eva_list, columns=['mae', 'rmse', 'mape', 'r2']).describe().to_csv('./eva_files/baseline_mask_loss.csv', encoding='utf-8-sig')

In [26]:

visualize_feature(data[5], masked_data[5], reconstructed[5], 'NO2')

---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
Cell In[26], line 1
----> 1 visualize_feature(data[5], masked_data[5], reconstructed[5], 'NO2')

NameError: name 'data' is not defined

In [ ]: