ASE-JSCCtrain.py

import torch
import argparse
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
# from torchvision.models import resnet18  # 以ResNet18为例，也可以根据实际情况选择其他模型
from torchvision.models import resnet18, ResNet18_Weights
import torch.nn.functional as F
from torch.utils.tensorboard import SummaryWriter
import time
from torch.optim.lr_scheduler import ReduceLROnPlateau
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

# device = torch.device("cpu") 

transform = transforms.Compose([
    transforms.Resize((256, 256)),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

mean = 0  
std_dev = 0.1  
train_dataset = datasets.ImageFolder(root='data/UCMerced_LandUse-train/Images', transform=transform)
test_dataset = datasets.ImageFolder(root='data/UCMerced_LandUse-test/Images', transform=transform)

train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False)


# The (real) AWGN channel    
def AWGN_channel(x, snr, P=2):
    batch_size, channels, height, width = x.shape  
    gamma = 10 ** (snr / 10.0)      
    noise = torch.sqrt(P / gamma) * torch.randn(batch_size, channels, height, width).to(device)  
    y = x + noise   
    return y


# Please set the symbol power if it is not a default value
def Fading_channel(x, snr, P = 2):
    gamma = 10 ** (snr / 10.0)
    [batch_size, feature_length] = x.shape
    K = feature_length//2
    
    h_I = torch.randn(batch_size, K).to(device)
    h_R = torch.randn(batch_size, K).to(device) 
    h_com = torch.complex(h_I, h_R)  
    x_com = torch.complex(x[:, 0:feature_length:2], x[:, 1:feature_length:2])
    y_com = h_com*x_com
    
    n_I = torch.sqrt(P/gamma)*torch.randn(batch_size, K).to(device)
    n_R = torch.sqrt(P/gamma)*torch.randn(batch_size, K).to(device)
    noise = torch.complex(n_I, n_R)
    
    y_add = y_com + noise
    y = y_add/h_com
    
    y_out = torch.zeros(batch_size, feature_length).to(device)
    y_out[:, 0:feature_length:2] = y.real
    y_out[:, 1:feature_length:2] = y.imag
    return y_out


def Combined_channel(x, snr, batch_size, channel, height, width):
    P=2
    x_faded = Fading_channel(x, snr, P)
    print ("x_faded.shape:",x_faded.shape)
    x_faded = x_faded.view((batch_size, channel, height, width))
    print ("x_faded.view.shape:",x_faded.shape)
    snr = torch.randint(0, 28, (x_faded.shape[0], x_faded.shape[1], x_faded.shape[2], 1)).to(device)
    x_combined = AWGN_channel(x_faded, snr, P)
    return x_combined


def Channel(z, snr, channel_type, batch_size, channel, height, width):
    if channel_type == 'AWGN':
        z = AWGN_channel(z, snr)
    elif channel_type == 'Fading':
        z = Fading_channel(z, snr)
    elif channel_type == 'Combined_channel':
        z = Combined_channel(z, snr, batch_size, channel, height, width)
    return z

class Autoencoder(nn.Module):
    def __init__(self):
        super(Autoencoder, self).__init__()
        self.encoder = nn.Sequential(
            nn.Conv2d(512, 256, kernel_size=3, stride=1, padding=1), 
            nn.ReLU(),
            nn.Conv2d(256, 128, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.Conv2d(128, 64, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.Conv2d(64, 32, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
        )
        self.flatten = nn.Flatten() 
        
        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(32, 64, kernel_size=3, stride=1, padding=1),  
            nn.ReLU(),
            nn.ConvTranspose2d(64, 128, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.ConvTranspose2d(128, 256, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.ConvTranspose2d(256, 512, kernel_size=3, stride=1, padding=1),
            nn.Sigmoid()
        )

    def forward(self, x, channel_type):
        # x = self.encoder(x)
        # print("encoder x.shape:",x.shape)
        # noise = torch.randn_like(x) * std_dev + mean
        # x = x + noise
        batch_size, channel, height, width = x.shape
        if channel_type == 'Fading' or channel_type == 'Combined_channel':
            x = self.flatten(x)
            print("after flatten x.shape", x.shape)
            SNR = torch.randint(0, 28, (x.shape[0], 1)).to(device)
        else :
            SNR = torch.randint(0, 28, (x.shape[0], x.shape[1], x.shape[2], 1)).to(device)
        x = Channel(x, SNR, channel_type, batch_size, channel, height, width)
        print("after Channel x.shape:",x.shape)
        x = x.view((batch_size, channel, height, width))     
        # x = self.decoder(x)
        return x



def mask_gen(weights, cr):
    position = round(cr*weights.size(1))
    weights_sorted, index = torch.sort(weights, dim=1)
    mask = torch.zeros_like(weights)

    for i in range(weights.size(0)):
        weight = weights_sorted[i, position-1]
        # print(weight)
        for j in range(weights.size(1)):
            if weights[i, j] <= weight:
                mask[i, j] = 1
    return mask  
  

class SE_Block(nn.Module):
    def __init__(self, inchannel, ratio=16):
        super(SE_Block, self).__init__()
        self.gap = nn.AdaptiveAvgPool2d((1,1))
        self.fc = nn.Sequential(
            nn.Linear(inchannel, inchannel // ratio, bias=False),  # 从 c -> c/r
            nn.ReLU(),
            nn.Linear(inchannel // ratio, inchannel, bias=False),  # 从 c/r -> c
            nn.Sigmoid()
        )

    def forward(self, x, cr=0.8):
        b, c, h, w = x.size()
        y = self.gap(x).view(b, c)
        print("y shape of Fsq:", y.shape)
        y = self.fc(y)
        print("y shape of Fex:", y.shape)
        mask = mask_gen(y, cr).view(b,c,1,1)
        print("mask shape:", mask.shape)
        print("x shape:", x.shape)
        return x * mask


class SatelliteClassifierWithAttention(nn.Module):
    def __init__(self, num_classes, autoencoder_out_channels=None):
        super(SatelliteClassifierWithAttention, self).__init__()
        
        # 1. 加载 ResNet18 (使用新的 weights 参数消除警告)
        weights = ResNet18_Weights.IMAGENET1K_V1
        self.backbone = resnet18(weights=weights)
        
        # 2. 移除原有的分类头 (fc) 和 全局平均池化 (avgpool)，我们要自己控制流程
        # 注意：我们需要保留 avgpool 的逻辑，但要在我们的模块之后执行，
        # 或者如果 Autoencoder 改变了空间维度，我们需要重新定义 pool。
        # 这里我们暂时移除原有的 fc，保留 backbone 到 layer4 的输出能力。
        self.backbone.fc = nn.Identity() 
        self.backbone.avgpool = nn.Identity() # 先移除，我们在后面手动加或根据情况处理

        # 获取 ResNet18 layer4 输出的通道数 (通常是 512)
        resnet_out_channels = self.backbone.layer4[-1].conv2.out_channels
        
        # 3. 定义自定义模块
        self.attention_module = SE_Block(resnet_out_channels)
        
        # 实例化 Autoencoder
        self.autoencoder = Autoencoder()
        
        # 【重要】获取 Autoencoder 的输出通道数
        # 方法 A: 如果 Autoencoder 有属性说明输出通道
        if hasattr(self.autoencoder, 'out_channels'):
            final_features = self.autoencoder.out_channels
        # 方法 B: 如果不知道，可以通过一次假推理获取 (推荐用于灵活架构)
        else:
            dummy_input = torch.zeros(1, resnet_out_channels, 7, 7) # 假设输入是 7x7 (ResNet18 标准输出)
            # 注意：这里需要知道 channel_type 和 cr 才能跑通假推理，如果复杂则需硬编码或修改 Autoencoder 设计
            # 为了安全，建议在 Autoencoder 类里明确写上 out_channels 属性
            final_features = resnet_out_channels # 默认 fallback，最好修改 Autoencoder 添加该属性

        # 4. 定义最终的分类器 (Global Average Pooling + Linear)
        self.global_pool = nn.AdaptiveAvgPool2d(1) # 适应任意 H, W
        self.classifier = nn.Linear(final_features, num_classes)

    def forward(self, x, cr, channel_type):
        # --- 阶段 1: ResNet 骨干特征提取 (直到 layer4) ---
        # 手动复现 ResNet 的前半部分，直到 layer4 输出
        x = self.backbone.conv1(x)
        x = self.backbone.bn1(x)
        x = self.backbone.relu(x)
        x = self.backbone.maxpool(x)
        
        x = self.backbone.layer1(x)
        x = self.backbone.layer2(x)
        x = self.backbone.layer3(x)
        x = self.backbone.layer4(x)
        
        # 此时 x shape: [B, 512, H, W] (通常 H=W=7 对于 224 输入)
        # print(f"After ResNet: {x.shape}")

        # --- 阶段 2: 注意力机制 ---
        # SE_Block 应该保持 [B, C, H, W] 不变，或者只改变 C
        x = self.attention_module(x, cr)
        # print(f"After Attention: {x.shape}")

        # --- 阶段 3: 自动编码器 (JSCC 核心) ---
        # Autoencoder 可能会改变通道数，也可能改变 H, W (如果是卷积下采样)
        x = self.autoencoder(x, channel_type)
        # print(f"After Autoencoder: {x.shape}")

        # --- 阶段 4: 全局池化与分类 ---
        # 使用 AdaptiveAvgPool2d 确保无论 H,W 是多少都能变成 1x1
        x = self.global_pool(x)
        
        # 展平 [B, C, 1, 1] -> [B, C]
        x = torch.flatten(x, 1)
        
        # 分类
        x = self.classifier(x)
        
        return x

def continue_train(cr, num_epochs, pre_checkpoint, channel_type):
    start_time = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())
    num_classes = len(train_dataset.classes)  
    model = SatelliteClassifierWithAttention(num_classes)
    model = model.to(device)

    pretrained_dict = torch.load(f'{pre_checkpoint}')
    model_dict = model.state_dict()

    pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
    model_dict.update(pretrained_dict)
    model.load_state_dict(model_dict)
    
    optimizer = optim.Adam(model.parameters(), lr=0.0001)
    scheduler = ReduceLROnPlateau(optimizer, mode='min', factor=0.1, patience=5)
    criterion = nn.CrossEntropyLoss()
    writer = SummaryWriter()

    # num_epochs = 50 

    for epoch in range(num_epochs):
        model.train()
        running_loss = 0.0

        for images, labels in train_loader:
            images, labels = images.to(device), labels.to(device)
            optimizer.zero_grad()
            outputs = model(images, cr, channel_type)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

            running_loss += loss.item()

        print(f'Epoch {epoch+1}/{num_epochs}, Loss: {running_loss/len(train_loader)}')
        avg_train_loss = running_loss / len(train_loader)
        scheduler.step(avg_train_loss)
        writer.add_scalar('Training Loss', running_loss/len(train_loader), epoch + 1)

        model.eval()
        correct = 0
        total = 0

        with torch.no_grad():
            for images, labels in test_loader:
                images, labels = images.to(device), labels.to(device)
                outputs = model(images, cr, channel_type)
                _, predicted = torch.max(outputs, 1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()

        accuracy = correct / total
        print(f'Test Accuracy: {accuracy}')

        writer.add_scalar('Test Accuracy', accuracy)
# Save the model with the specified cr and num_epochs in the file name
    save_path = f'checkpoint/classifier_attention_auto_UCMerced_LandUse_{channel_type}_ResNet18_60epoch_0.5_up_{num_epochs}epoch_{cr}.pth'
    torch.save(model.state_dict(), save_path)

    writer.close()

    current_time = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())
    # Write results to a txt file
    with open(f'logs/ResNet18/classifier_attention_auto_UCMerced_LandUse_{channel_type}_ResNet18_60_up_{num_epochs}epoch_{cr}.txt', 'w') as file:
        file.write('strat comtinue training...\n')
        file.write(f'Time: {start_time}----------{current_time}\n')
        file.write(f'model name:{save_path}\n')
        file.write(f'channel_type:{channel_type}\n')
        file.write(f'CR (Compression Ratio): {cr}\n')
        file.write(f'Num Epochs: {num_epochs}\n')
        file.write(f'Test Accuracy: {accuracy}\n')
        file.write('train over!\n')

def train(cr, num_epochs, channel_type):
    start_time = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())
    num_classes = len(train_dataset.classes)  
    model = SatelliteClassifierWithAttention(num_classes)
    model = model.to(device)
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    scheduler = ReduceLROnPlateau(optimizer, mode='min', factor=0.1, patience=5)

    writer = SummaryWriter()
    # num_epochs = 20

    for epoch in range(num_epochs):
        model.train()
        running_loss = 0.0

        for images, labels in train_loader:
            print("epoch:",epoch)
            print(images.shape)
            images, labels = images.to(device), labels.to(device)
            optimizer.zero_grad()
            outputs = model(images, cr, channel_type)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

            running_loss += loss.item()

        print(f'Epoch {epoch+1}/{num_epochs}, Loss: {running_loss/len(train_loader)}')
 
        avg_train_loss = running_loss / len(train_loader)

        scheduler.step(avg_train_loss)
        writer.add_scalar('Training Loss', running_loss/len(train_loader), epoch + 1)

        model.eval()
        correct = 0
        total = 0

        with torch.no_grad():
            for images, labels in test_loader:
                images, labels = images.to(device), labels.to(device)
                outputs = model(images, cr, channel_type)
                _, predicted = torch.max(outputs, 1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()

        accuracy = correct / total
        print(f'Test Accuracy: {accuracy}')
        writer.add_scalar('Test Accuracy', accuracy)

    # Save the model with the specified cr and num_epochs in the file name
    save_path = f'checkpoint/classifier_attention_auto_UCMerced_LandUse_Combined_channel_ResNet18_{num_epochs}epoch_{cr}.pth'
    torch.save(model.state_dict(), save_path)

    
    writer.close()
    current_time = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())
    # Write results to a txt file
    with open(f'logs/ResNet18/classifier_attention_auto_UCMerced_LandUse_Combined_channel_ResNet18_{num_epochs}epoch_{cr}.txt', 'w') as file:
        file.write('strat training...\n')
        file.write(f'Time: {start_time}----------{current_time}\n')
        file.write(f'model name:{save_path}\n')
        file.write(f'model name:{channel_type}\n')
        file.write(f'CR (Compression Ratio): {cr}\n')
        file.write(f'Num Epochs: {num_epochs}\n')
        file.write(f'Test Accuracy: {accuracy}\n')
        file.write('train over!\n')

def main(task,cr,num_epochs,pre_checkpoint,channel_type):
    if task == 'continue':
        print("continue_train start!")
        continue_train(cr, num_epochs,pre_checkpoint,channel_type)
        print("continue_train over!")
    else :
        print("train start!")
        train(cr, num_epochs,channel_type)
        print("train over!")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Train or continue training a model.")
    parser.add_argument('--task', choices=['continue', 'train'], default='train', required=True, help='Specify the task (continue or train).')
    parser.add_argument('--cr', type=float, default=0.1, help='Specify the compression ratio (cr) for the SE Block.')
    parser.add_argument('--num_epochs', type=int, default=60, help='Specify the number of epochs for training.')
    parser.add_argument('--pre_checkpoint', type=str, default=None, help='Specify the pretrained checkpoint for continue train.')
    parser.add_argument('--channel_type', choices=['AWGN', 'Fading',"Combined_channel"], default='Combined_channel', help='Specify the channel_type for transfer.')
    args = parser.parse_args()
    main(args.task, args.cr, args.num_epochs,args.pre_checkpoint,args.channel_type)