MAE_ATMO/torch_MAE_1d_final_20.ipynb

In [1]:

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 matplotlib.pyplot as plt
import cv2
import pandas as pd

In [2]:

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

Out[2]:

<torch._C.Generator at 0x7f8d8a9ef7f0>

In [3]:

# 计算图像数据中的最大像素值
max_pixel_value = 107.49169921875

In [4]:

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 [5]:

train_set = NO2Dataset(image_dir, mask_dir)
train_loader = DataLoader(train_set, 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 [6]:

# 可视化特定特征的函数
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 [7]:

class Conv(nn.Sequential):
    def __init__(self, in_channels, out_channels, kernel_size=3, dilation=1, stride=1, bias=False):
        super(Conv, self).__init__(
            nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, bias=bias,
                      dilation=dilation, stride=stride, padding=((stride - 1) + dilation * (kernel_size - 1)) // 2)
        )

In [8]:

class ConvBNReLU(nn.Sequential):
    def __init__(self, in_channels, out_channels, kernel_size=3, dilation=1, stride=1, norm_layer=nn.BatchNorm2d,
                 bias=False):
        super(ConvBNReLU, self).__init__(
            nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, bias=bias,
                      dilation=dilation, stride=stride, padding=((stride - 1) + dilation * (kernel_size - 1)) // 2),
            norm_layer(out_channels),
            nn.ReLU()
        )

In [9]:

class SeparableBNReLU(nn.Sequential):
    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, dilation=1, norm_layer=nn.BatchNorm2d):
        super(SeparableBNReLU, self).__init__(
            nn.Conv2d(in_channels, in_channels, kernel_size=kernel_size, stride=stride, dilation=dilation,
                      padding=((stride - 1) + dilation * (kernel_size - 1)) // 2, groups=in_channels, bias=False),
            # 分离卷积，仅调整空间信息
            norm_layer(in_channels),  # 对输入通道进行归一化
            nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False),  # 这里进行升维操作
            nn.ReLU6()
        )

In [10]:

class ResidualBlock(nn.Module):
    def __init__(self, in_channels, out_channels, stride=1, downsample=None):
        super(ResidualBlock, self).__init__()
        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(out_channels)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(out_channels)

        # 如果输入和输出通道不一致，进行降采样操作
        self.downsample = downsample
        if in_channels != out_channels or stride != 1:
            self.downsample = nn.Sequential(
                nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm2d(out_channels)
            )

    def forward(self, x):
        identity = x
        if self.downsample is not None:
            identity = self.downsample(x)

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)

        out += identity
        out = self.relu(out)
        return out

In [11]:

class Mlp(nn.Module):
    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.ReLU6, drop=0.):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        self.fc1 = nn.Conv2d(in_features, hidden_features, 1, 1, 0, bias=True)

        self.act = act_layer()
        self.fc2 = nn.Conv2d(hidden_features, out_features, 1, 1, 0, bias=True)
        self.drop = nn.Dropout(drop, inplace=True)

    def forward(self, x):
        x = self.fc1(x)
        x = self.act(x)
        x = self.drop(x)
        x = self.fc2(x)
        x = self.drop(x)
        return x

In [12]:

class MultiHeadAttentionBlock(nn.Module):
    def __init__(self, embed_dim, num_heads, dropout=0.1):
        super(MultiHeadAttentionBlock, self).__init__()
        self.attention = nn.MultiheadAttention(embed_dim, num_heads, dropout=dropout)
        self.norm = nn.LayerNorm(embed_dim)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        # (B, C, H, W) -> (HW, B, C) for MultiheadAttention compatibility
        B, C, H, W = x.shape
        x = x.view(B, C, H * W).permute(2, 0, 1)  # (B, C, H, W) -> (HW, B, C)

        # Apply multihead attention
        attn_output, _ = self.attention(x, x, x)

        # Apply normalization and dropout
        attn_output = self.norm(attn_output)
        attn_output = self.dropout(attn_output)

        # Reshape back to (B, C, H, W)
        attn_output = attn_output.permute(1, 2, 0).view(B, C, H, W)

        return attn_output

In [13]:

class SpatialAttentionBlock(nn.Module):
    def __init__(self):
        super(SpatialAttentionBlock, self).__init__()
        self.conv = nn.Conv2d(2, 1, kernel_size=7, padding=3, bias=False)

    def forward(self, x): #(B, 64, H, W)
        avg_out = torch.mean(x, dim=1, keepdim=True) #(B, 1, H, W)
        max_out, _ = torch.max(x, dim=1, keepdim=True)#(B, 1, H, W)
        out = torch.cat([avg_out, max_out], dim=1)#(B, 2, H, W)
        out = torch.sigmoid(self.conv(out))#(B, 1, H, W)
        return x * out #(B, C, H, W)

In [14]:

class DecoderAttentionBlock(nn.Module):
    def __init__(self, in_channels):
        super(DecoderAttentionBlock, self).__init__()
        self.conv1 = nn.Conv2d(in_channels, in_channels // 2, kernel_size=1)
        self.conv2 = nn.Conv2d(in_channels // 2, in_channels, kernel_size=1)
        self.spatial_attention = SpatialAttentionBlock()

    def forward(self, x):
        # 通道注意力
        b, c, h, w = x.size()
        avg_pool = F.adaptive_avg_pool2d(x, 1)
        max_pool = F.adaptive_max_pool2d(x, 1)

        avg_out = self.conv1(avg_pool)
        max_out = self.conv1(max_pool)

        out = avg_out + max_out
        out = torch.sigmoid(self.conv2(out))

        # 添加空间注意力
        out = x * out
        out = self.spatial_attention(out)
        return out

In [15]:

class SEBlock(nn.Module):
    def __init__(self, in_channels, reduced_dim):
        super(SEBlock, self).__init__()
        self.se = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),  # 全局平均池化
            nn.Conv2d(in_channels, reduced_dim, kernel_size=1),
            nn.ReLU(),
            nn.Conv2d(reduced_dim, in_channels, kernel_size=1),
            nn.Sigmoid()  # 使用Sigmoid是因为我们要对通道进行权重归一化
        )

    def forward(self, x):
        return x * self.se(x)

In [16]:

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

In [17]:

# 定义Masked Autoencoder模型
class MaskedAutoencoder(nn.Module):
    def __init__(self):
        super(MaskedAutoencoder, self).__init__()
        self.encoder = nn.Sequential(
            Conv(1, 32, kernel_size=3, stride=2),
            
            nn.ReLU(),
            
            SEBlock(32,32),
            
            ConvBNReLU(32, 64, kernel_size=3, stride=2),
            
            ResidualBlock(64,64),
            
            SeparableBNReLU(64, 128, kernel_size=3, stride=2),
            
            MultiHeadAttentionBlock(embed_dim=128, num_heads=4),
            
            SEBlock(128, 128)
            
        )
        # self.mlp = Mlp(in_features=128, hidden_features=256, out_features=128, act_layer=nn.ReLU6, drop=0.1)
        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(128, 32, kernel_size=3, stride=2, padding=1, output_padding=1),
            nn.ReLU(),
            
            DecoderAttentionBlock(32),
            nn.ConvTranspose2d(32, 16, kernel_size=3, stride=2, padding=1, output_padding=1),
            nn.ReLU(),
            
            DecoderAttentionBlock(16),
            nn.ReLU(),
            
            nn.ConvTranspose2d(16, 1, kernel_size=3, stride=2, padding=1, output_padding=1),  # 修改为 output_padding=1
            nn.Sigmoid()
        )

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

# 实例化模型、损失函数和优化器
model = MaskedAutoencoder()
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.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)
            running_loss += loss.item()
    return running_loss / (batch_idx + 1)

In [20]:

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

cuda

In [21]:

model = model.to(device)

num_epochs = 160
train_losses = list()
val_losses = list()
for epoch in range(num_epochs):
    train_loss = train_epoch(model, device, train_loader, 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}')

/root/miniconda3/envs/python38/lib/python3.8/site-packages/torch/nn/modules/conv.py:456: UserWarning: Applied workaround for CuDNN issue, install nvrtc.so (Triggered internally at /opt/conda/conda-bld/pytorch_1711403590347/work/aten/src/ATen/native/cudnn/Conv_v8.cpp:80.)
  return F.conv2d(input, weight, bias, self.stride,

Epoch 1, Train Loss: 1.828955806274876, Val Loss: 0.08777590596408986
Epoch 2, Train Loss: 0.06457909727781012, Val Loss: 0.05018303115198861
Epoch 3, Train Loss: 0.04399169035006368, Val Loss: 0.03933813378437242
Epoch 4, Train Loss: 0.03737294341049839, Val Loss: 0.04090026577017201
Epoch 5, Train Loss: 0.03340746862947513, Val Loss: 0.029788545930563515
Epoch 6, Train Loss: 0.03127880183240158, Val Loss: 0.02878953230136366
Epoch 7, Train Loss: 0.030086695853816837, Val Loss: 0.027378849156979305
Epoch 8, Train Loss: 0.02827827861470184, Val Loss: 0.026564865748384105
Epoch 9, Train Loss: 0.026973650764014447, Val Loss: 0.026876062349374615
Epoch 10, Train Loss: 0.026198443756149145, Val Loss: 0.025235873994542593
Epoch 11, Train Loss: 0.025248640154501754, Val Loss: 0.025164278752323407
Epoch 12, Train Loss: 0.0246738152373493, Val Loss: 0.02402887423870279
Epoch 13, Train Loss: 0.02429686849446673, Val Loss: 0.02467221769490349
Epoch 14, Train Loss: 0.023617587716242915, Val Loss: 0.024100169289245535
Epoch 15, Train Loss: 0.022902602209535796, Val Loss: 0.023378314227977797
Epoch 16, Train Loss: 0.022661644239067746, Val Loss: 0.02472560463556603
Epoch 17, Train Loss: 0.02193861959154526, Val Loss: 0.02273730694580434
Epoch 18, Train Loss: 0.021775715561075645, Val Loss: 0.022977211248518814
Epoch 19, Train Loss: 0.021564541914852325, Val Loss: 0.022313175500551268
Epoch 20, Train Loss: 0.0214472935851396, Val Loss: 0.022048505606935984
Epoch 21, Train Loss: 0.020810687219340835, Val Loss: 0.02184077285563768
Epoch 22, Train Loss: 0.020310772384592647, Val Loss: 0.021513454977478554
Epoch 23, Train Loss: 0.02010334756350118, Val Loss: 0.02177375905326943
Epoch 24, Train Loss: 0.02025744297795675, Val Loss: 0.02049418441506464
Epoch 25, Train Loss: 0.019826160295995657, Val Loss: 0.023377947564890134
Epoch 26, Train Loss: 0.019065276574806875, Val Loss: 0.020193443425110917
Epoch 27, Train Loss: 0.01881279432745071, Val Loss: 0.01942526154331307
Epoch 28, Train Loss: 0.01839842515413841, Val Loss: 0.01973166508572315
Epoch 29, Train Loss: 0.018092166516555555, Val Loss: 0.021518220902601286
Epoch 30, Train Loss: 0.01789530134942543, Val Loss: 0.0191833000741343
Epoch 31, Train Loss: 0.017643442852021546, Val Loss: 0.018857373494599292
Epoch 32, Train Loss: 0.017585936365604543, Val Loss: 0.018622038858150367
Epoch 33, Train Loss: 0.017121152348513382, Val Loss: 0.018597172726112516
Epoch 34, Train Loss: 0.016807572604223872, Val Loss: 0.01907729919054615
Epoch 35, Train Loss: 0.0167503119735983, Val Loss: 0.018055098590010137
Epoch 36, Train Loss: 0.01674377040839509, Val Loss: 0.017786314029858183
Epoch 37, Train Loss: 0.016270555827641888, Val Loss: 0.01821137344770467
Epoch 38, Train Loss: 0.016271821564090166, Val Loss: 0.017419732745681236
Epoch 39, Train Loss: 0.01634730132180823, Val Loss: 0.017153916838787385
Epoch 40, Train Loss: 0.016149515664855545, Val Loss: 0.01720947952968861
Epoch 41, Train Loss: 0.015722640304331573, Val Loss: 0.01671495117636314
Epoch 42, Train Loss: 0.015584125958882165, Val Loss: 0.016605446490445243
Epoch 43, Train Loss: 0.015607581132996168, Val Loss: 0.016551834531128407
Epoch 44, Train Loss: 0.015686789721375303, Val Loss: 0.017196020681355426
Epoch 45, Train Loss: 0.0152399734099302, Val Loss: 0.016840887422770706
Epoch 46, Train Loss: 0.015122933551651296, Val Loss: 0.018965846010998114
Epoch 47, Train Loss: 0.015065566115259554, Val Loss: 0.016344470375064594
Epoch 48, Train Loss: 0.014854169773766726, Val Loss: 0.016327281677122437
Epoch 49, Train Loss: 0.014882152102459845, Val Loss: 0.015837757153186336
Epoch 50, Train Loss: 0.014656414190957848, Val Loss: 0.016042638750774645
Epoch 51, Train Loss: 0.014637816764200418, Val Loss: 0.015558397091591536
Epoch 52, Train Loss: 0.01454300198784214, Val Loss: 0.015685647628756603
Epoch 53, Train Loss: 0.014566657712691994, Val Loss: 0.01571561763090874
Epoch 54, Train Loss: 0.01434676954522729, Val Loss: 0.015356795890117758
Epoch 55, Train Loss: 0.014364799384348557, Val Loss: 0.015472657116713808
Epoch 56, Train Loss: 0.014128341450930783, Val Loss: 0.015367844809235922
Epoch 57, Train Loss: 0.014267995692878677, Val Loss: 0.016404178910958234
Epoch 58, Train Loss: 0.01399662052882773, Val Loss: 0.014956932640008962
Epoch 59, Train Loss: 0.013984658806607056, Val Loss: 0.01512009026343698
Epoch 60, Train Loss: 0.013917681792278608, Val Loss: 0.01516334629103319
Epoch 61, Train Loss: 0.013808810461811614, Val Loss: 0.015075811351746765
Epoch 62, Train Loss: 0.014042920544387051, Val Loss: 0.015152243647112776
Epoch 63, Train Loss: 0.0136711714971971, Val Loss: 0.014804388201837219
Epoch 64, Train Loss: 0.013782783121797457, Val Loss: 0.015533475858618074
Epoch 65, Train Loss: 0.013631306383669661, Val Loss: 0.014752479089396213
Epoch 66, Train Loss: 0.013644688259186357, Val Loss: 0.01469478735338841
Epoch 67, Train Loss: 0.013522711930056793, Val Loss: 0.014726998854372928
Epoch 68, Train Loss: 0.01350348583159692, Val Loss: 0.014617940202466588
Epoch 69, Train Loss: 0.013397794087644684, Val Loss: 0.014498871904033334
Epoch 70, Train Loss: 0.013320690925504888, Val Loss: 0.014324163573224153
Epoch 71, Train Loss: 0.013295841332008108, Val Loss: 0.014810262790033177
Epoch 72, Train Loss: 0.013151036726943614, Val Loss: 0.014535954208182754
Epoch 73, Train Loss: 0.01315474125409597, Val Loss: 0.014322022976937578
Epoch 74, Train Loss: 0.013201014497473337, Val Loss: 0.014625799591972757
Epoch 75, Train Loss: 0.013166735187155065, Val Loss: 0.01410402478511209
Epoch 76, Train Loss: 0.013011173492199496, Val Loss: 0.014279130234647153
Epoch 77, Train Loss: 0.012954122741131833, Val Loss: 0.015670507896079947
Epoch 78, Train Loss: 0.012964830874202497, Val Loss: 0.013965579806201493
Epoch 79, Train Loss: 0.01284469154765874, Val Loss: 0.014020084167149529
Epoch 80, Train Loss: 0.01269332727230194, Val Loss: 0.014467649356420361
Epoch 81, Train Loss: 0.012900225120779287, Val Loss: 0.014321781124975255
Epoch 82, Train Loss: 0.012758908171705795, Val Loss: 0.013745425046602292
Epoch 83, Train Loss: 0.01266205709418683, Val Loss: 0.013802579048075784
Epoch 84, Train Loss: 0.012549680232128315, Val Loss: 0.013783436657777473
Epoch 85, Train Loss: 0.012634162601689545, Val Loss: 0.01444499020867828
Epoch 86, Train Loss: 0.012543465024190086, Val Loss: 0.014219797327558495
Epoch 87, Train Loss: 0.012490486795234195, Val Loss: 0.013482047425610806
Epoch 88, Train Loss: 0.012537837625619327, Val Loss: 0.014496686354057113
Epoch 89, Train Loss: 0.012536356080786891, Val Loss: 0.013949389360956292
Epoch 90, Train Loss: 0.012426643302601776, Val Loss: 0.013645224328806152
Epoch 91, Train Loss: 0.012394862496806531, Val Loss: 0.013617335818707943
Epoch 92, Train Loss: 0.012383774110075959, Val Loss: 0.013630805342499889
Epoch 93, Train Loss: 0.012307288521749267, Val Loss: 0.013647960637932393
Epoch 94, Train Loss: 0.012298794681625218, Val Loss: 0.013733426678870151
Epoch 95, Train Loss: 0.012473734824263165, Val Loss: 0.013764488983398942
Epoch 96, Train Loss: 0.012222074678515276, Val Loss: 0.013446863671180918
Epoch 97, Train Loss: 0.012306330008120344, Val Loss: 0.013694896279319899
Epoch 98, Train Loss: 0.012166704374263019, Val Loss: 0.013338639831809855
Epoch 99, Train Loss: 0.012187617220447965, Val Loss: 0.01352898025913025
Epoch 100, Train Loss: 0.012234464256565252, Val Loss: 0.013427354033980796
Epoch 101, Train Loss: 0.012252488267122273, Val Loss: 0.013189904238861887
Epoch 102, Train Loss: 0.01208857831692225, Val Loss: 0.013358786896760785
Epoch 103, Train Loss: 0.012067412587693718, Val Loss: 0.013412703287356826
Epoch 104, Train Loss: 0.011943526178348863, Val Loss: 0.013329273687480991
Epoch 105, Train Loss: 0.012186939030457911, Val Loss: 0.013039200052396576
Epoch 106, Train Loss: 0.012064487648833739, Val Loss: 0.013328265718448518
Epoch 107, Train Loss: 0.01196315302624942, Val Loss: 0.013011285284561898
Epoch 108, Train Loss: 0.011942964125175082, Val Loss: 0.013228343076892753
Epoch 109, Train Loss: 0.011851983095862363, Val Loss: 0.012941466032791494
Epoch 110, Train Loss: 0.011892807039401035, Val Loss: 0.013264400856708413
Epoch 111, Train Loss: 0.011915889784747192, Val Loss: 0.01319889353115612
Epoch 112, Train Loss: 0.011905829402123484, Val Loss: 0.014149442662610047
Epoch 113, Train Loss: 0.011818570989455903, Val Loss: 0.013042371636673586
Epoch 114, Train Loss: 0.011752497955140743, Val Loss: 0.01301327784226012
Epoch 115, Train Loss: 0.011813209191606375, Val Loss: 0.01286677592225484
Epoch 116, Train Loss: 0.011725439075113198, Val Loss: 0.013167357391941904
Epoch 117, Train Loss: 0.011835235226721141, Val Loss: 0.01286814648157625
Epoch 118, Train Loss: 0.011680879099873835, Val Loss: 0.012708428107313256
Epoch 119, Train Loss: 0.01173722647959322, Val Loss: 0.012885383775096331
Epoch 120, Train Loss: 0.011672099965343777, Val Loss: 0.012913884747940214
Epoch 121, Train Loss: 0.011704605972866693, Val Loss: 0.012728425813143823
Epoch 122, Train Loss: 0.011705320578015021, Val Loss: 0.012817327530860012
Epoch 123, Train Loss: 0.011644495068492288, Val Loss: 0.012942980015789396
Epoch 124, Train Loss: 0.011633442955439171, Val Loss: 0.012936850551015405
Epoch 125, Train Loss: 0.011616052921558396, Val Loss: 0.012702107387803384
Epoch 126, Train Loss: 0.011607619160652588, Val Loss: 0.012658866025062639
Epoch 127, Train Loss: 0.011635440495310788, Val Loss: 0.01304104494681554
Epoch 128, Train Loss: 0.01150463111074775, Val Loss: 0.013212839975508291
Epoch 129, Train Loss: 0.011585681133293078, Val Loss: 0.01278914052492647
Epoch 130, Train Loss: 0.011392400087565896, Val Loss: 0.012796499154794572
Epoch 131, Train Loss: 0.011433751801358598, Val Loss: 0.012598757076063264
Epoch 132, Train Loss: 0.011496097840921303, Val Loss: 0.01271620902941743
Epoch 133, Train Loss: 0.011477598884815804, Val Loss: 0.013398304248034065
Epoch 134, Train Loss: 0.011365674946314552, Val Loss: 0.012668505741922713
Epoch 135, Train Loss: 0.01142354957696995, Val Loss: 0.013356663286685944
Epoch 136, Train Loss: 0.011355750374139497, Val Loss: 0.012617305616167054
Epoch 137, Train Loss: 0.011350866257877013, Val Loss: 0.012997348792850971
Epoch 138, Train Loss: 0.011416472670617715, Val Loss: 0.012524361819473665
Epoch 139, Train Loss: 0.011427981736646458, Val Loss: 0.012654973694415235
Epoch 140, Train Loss: 0.011318818902213607, Val Loss: 0.012664613897787102
Epoch 141, Train Loss: 0.011320005095247446, Val Loss: 0.012727182441905363
Epoch 142, Train Loss: 0.011245375826651827, Val Loss: 0.012474427931010723
Epoch 143, Train Loss: 0.011338526420919091, Val Loss: 0.012642348824597117
Epoch 144, Train Loss: 0.011243535689207497, Val Loss: 0.012692421772030752
Epoch 145, Train Loss: 0.011166462189023289, Val Loss: 0.01263011310861182
Epoch 146, Train Loss: 0.011227301243942178, Val Loss: 0.012461379587427894
Epoch 147, Train Loss: 0.01119774208364019, Val Loss: 0.012749987918494353
Epoch 148, Train Loss: 0.011138954723441001, Val Loss: 0.012676928915194612
Epoch 149, Train Loss: 0.011145075226122398, Val Loss: 0.012806226499378681
Epoch 150, Train Loss: 0.011238663441737731, Val Loss: 0.012608930385157244
Epoch 151, Train Loss: 0.01112103075430724, Val Loss: 0.012799791727604261
Epoch 152, Train Loss: 0.01109027168958595, Val Loss: 0.01240885794273953
Epoch 153, Train Loss: 0.011098397055721026, Val Loss: 0.012326594039019364
Epoch 154, Train Loss: 0.011026590389676356, Val Loss: 0.012310143629672811
Epoch 155, Train Loss: 0.011067607804339682, Val Loss: 0.01242478439278567
Epoch 156, Train Loss: 0.01105262930215332, Val Loss: 0.01238662200465576
Epoch 157, Train Loss: 0.010977347388097117, Val Loss: 0.012163419262575569
Epoch 158, Train Loss: 0.010957017552071924, Val Loss: 0.012397716572480415
Epoch 159, Train Loss: 0.010956506543396192, Val Loss: 0.012370292931350309
Epoch 160, Train Loss: 0.01093887382980133, Val Loss: 0.012291266110294791
Test Loss: 0.006885056002065539

In [23]:

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

Out[23]:

<matplotlib.legend.Legend at 0x7f8d2cf4ebe0>

In [24]:

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

In [25]:

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 [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)
        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)
        ioa = cal_ioa(data_label.detach().numpy(), recon_no2.detach().numpy())
        r = np.corrcoef(data_label, recon_no2)[0, 1]
        eva_list.append([mae, rmse, mape, r2, ioa, r])

In [27]:

eva_list_frame = list()
device = 'cpu'
model = model.to(device)
best_mape = 1
best_img = None
best_mask = None
best_recov = None
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])
            if mape < best_mape:
                best_recov = rev_recon[i][0].numpy()
                best_mask = used_mask.numpy()
                best_img = sample[0].numpy()
                best_mape = mape

In [28]:

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

Out[28]:

	mae	rmse	mape	r2	ioa	r
count	4739.000000	4739.000000	4739.000000	4739.000000	4739.000000	4739.000000
mean	1.261634	1.801726	0.153962	0.681159	0.891040	0.840609
std	0.572205	0.861009	0.065723	0.249771	0.110411	0.124012
min	0.361480	0.468918	0.047540	-2.107971	-0.424296	-0.070884
25%	0.828453	1.149391	0.111256	0.600440	0.868937	0.797875
50%	1.135805	1.621294	0.143929	0.740937	0.922953	0.872734
75%	1.557381	2.250718	0.179544	0.835907	0.953556	0.921983
max	5.733449	8.356097	1.116946	0.985570	0.996237	0.993398

In [31]:

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

Out[31]:

	mae	rmse	mape	r2	ioa	r
count	75.000000	75.000000	75.000000	75.000000	75.000000	75.000000
mean	1.263991	1.987788	0.153931	0.907729	0.974785	0.953238
std	0.108035	0.209185	0.007592	0.017280	0.005782	0.007909
min	1.077143	1.658797	0.135271	0.791607	0.933031	0.905484
25%	1.208991	1.892923	0.149006	0.901544	0.972912	0.950092
50%	1.255151	1.967265	0.153929	0.908183	0.974939	0.953771
75%	1.307615	2.079039	0.158463	0.915666	0.977269	0.957454
max	1.956845	3.320712	0.175028	0.931715	0.981832	0.965467

In [67]:

# torch.save(model, './models/MAE/final_20.pt')

In [32]:

model_20 = torch.load('./models/MAE/final_20.pt')

In [38]:

# 可视化特定特征的函数
def visualize_rst(input_feature,masked_feature, recov_region, output_feature, title):
    plt.figure(figsize=(12, 6))
    plt.subplot(1, 4, 1)
    plt.imshow(input_feature, cmap='RdYlGn_r')
    plt.gca().axis('off')  # 获取当前坐标轴并关闭
    
    plt.subplot(1, 4, 2)
    plt.imshow(masked_feature, cmap='gray')
    plt.gca().axis('off')  # 获取当前坐标轴并关闭
    plt.subplot(1, 4, 3)
    plt.imshow(recov_region, cmap='RdYlGn_r')
    plt.gca().axis('off')  # 获取当前坐标轴并关闭
    plt.subplot(1, 4, 4)
    plt.imshow(output_feature, cmap='RdYlGn_r')
    plt.gca().axis('off')  # 获取当前坐标轴并关闭
    plt.savefig('./figures/result/20_samples.png', bbox_inches='tight')

In [39]:

best_mask_cp = np.where(best_mask == 0, np.nan, best_mask)

In [49]:

visualize_rst(best_img, best_mask, best_recov*best_mask_cp, best_img * (1-best_mask) + best_recov*best_mask, '')

In [33]:

find_ex = set([x.split('-')[0].strip() for x in os.listdir('./test_img/') if 'npy' in x])
find_ex

Out[33]:

{'1114', '1952', '2568', '3523', '602'}

In [70]:

for j in find_ex:
    ori = np.load(f'./test_img/{j}-real.npy')[0]
    mask = np.load(f'./test_img/{j}-mask.npy')
    mask_rev = 1 - mask
    img_in = ori * mask_rev / max_pixel_value
    img_out = model(torch.tensor(img_in.reshape(1, 1, 96, 96), dtype=torch.float32)).detach().cpu().numpy()[0][0] * max_pixel_value
    out = ori * mask_rev + img_out * mask
    plt.imshow(out, cmap='RdYlGn_r')
    plt.gca().axis('off')
    plt.savefig(f'./test_img/out_fig/{j}-mae_my_out.png', bbox_inches='tight')
    plt.clf()

<Figure size 640x480 with 0 Axes>

In [ ]: