AL for CNN

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader, Subset
from torchvision import transforms, models
import pandas as pd
import os
from PIL import Image
from sklearn.model_selection import train_test_split
import numpy as np
from time import time

start_time = time()


# 定义数据集类
class ChestXRayDataset(Dataset):
    def __init__(self, csv_path, root_dir, transform=None):
        self.data_info = pd.read_csv(csv_path)
        self.root_dir = root_dir
        self.transform = transform

    def __len__(self):
        return len(self.data_info)

    def __getitem__(self, idx):
        img_name = os.path.join(self.root_dir, self.data_info.iloc[idx, 0])
        image = Image.open(img_name).convert('RGB')
        label = self.data_info.iloc[idx, 1]
        if label == 'covid':
            label = 0
        elif label == 'normal':
            label = 1
        else:
            label = 2

        if self.transform:
            image = self.transform(image)

        return image, label

# 数据预处理
data_transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# 加载数据
root_dir = r'E:\NiuCode\LianXi\data\Chest X-ray\DataSet'
csv_path = r'E:\NiuCode\LianXi\data\Chest X-ray\metadata.csv'
dataset = ChestXRayDataset(csv_path, root_dir, transform=data_transform)

# 划分训练集、无标记样本池和测试集
train_data, test_data = train_test_split(dataset, test_size=0.1, random_state=42)
# 初始训练集只包含少量已标记数据
initial_train_size = 32
train_indices = np.random.choice(len(train_data), initial_train_size, replace=False)
train_subset = Subset(train_data, train_indices)
unlabeled_indices = [i for i in range(len(train_data)) if i not in train_indices]
unlabeled_pool = Subset(train_data, unlabeled_indices)

# 创建数据加载器
train_loader = DataLoader(train_subset, batch_size=32, shuffle=True)
test_loader = DataLoader(test_data, batch_size=32, shuffle=False)

# 定义CNN模型
class CNNModel(nn.Module):
    def __init__(self):
        super(CNNModel, self).__init__()
        self.conv1 = nn.Conv2d(3, 16, kernel_size=3, padding=1)
        self.relu1 = nn.ReLU()
        self.pool1 = nn.MaxPool2d(2, 2)

        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, padding=1)
        self.relu2 = nn.ReLU()
        self.pool2 = nn.MaxPool2d(2, 2)

        self.fc1 = nn.Linear(32 * 56 * 56, 128)
        self.relu3 = nn.ReLU()
        self.fc2 = nn.Linear(128, 3)

    def forward(self, x):
        out = self.pool1(self.relu1(self.conv1(x)))
        out = self.pool2(self.relu2(self.conv2(out)))
        out = out.view(-1, 32 * 56 * 56)
        out = self.relu3(self.fc1(out))
        out = self.fc2(out)
        return out

# 标记预算
labeling_budget = 640
batch_size = 32
num_iterations = labeling_budget // batch_size

# 检查 CUDA 是否可用
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

for iteration in range(num_iterations):
    model = CNNModel()
    # 将模型移到设备上
    model.to(device)

    # 定义损失函数和优化器
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=0.001)

    # 训练模型
    num_epochs = 10
    for epoch in range(num_epochs):
        running_loss = 0.0
        for i, (images, labels) in enumerate(train_loader):
            # 将数据移到设备上
            images = images.to(device)
            labels = labels.to(device)

            optimizer.zero_grad()
            outputs = model(images)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

            running_loss += loss.item()
        print(f'Epoch {epoch + 1}, Loss: {running_loss / len(train_loader)}')

    # 测试模型
    correct = 0
    total = 0
    with torch.no_grad():
        for images, labels in test_loader:
            # 将数据移到设备上
            images = images.to(device)
            labels = labels.to(device)

            outputs = model(images)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    test_accuracy = 100 * correct / total
    print(f'Iteration {iteration + 1}, Test Accuracy: {test_accuracy}%')

    # 基于熵最大化的不确定性采样
    unlabeled_loader = DataLoader(unlabeled_pool, batch_size=32, shuffle=False)
    all_entropies = []
    all_indices = []
    with torch.no_grad():
        for i, (images, _) in enumerate(unlabeled_loader):
            # 将数据移到设备上
            images = images.to(device)

            outputs = model(images)
            probabilities = torch.softmax(outputs, dim=1)
            entropies = -torch.sum(probabilities * torch.log(probabilities + 1e-10), dim=1)
            start_index = i * 32
            indices = list(range(start_index, start_index + len(images)))
            all_entropies.extend(entropies.cpu().numpy())
            all_indices.extend(indices)

    sorted_indices = np.argsort(all_entropies)[::-1]
    selected_indices = sorted_indices[:batch_size]
    selected_unlabeled_indices = [unlabeled_indices[i] for i in selected_indices]

    # 将选中的样本加入训练集
    new_train_indices = list(train_subset.indices) + selected_unlabeled_indices
    train_subset = Subset(train_data, new_train_indices)
    train_loader = DataLoader(train_subset, batch_size=32, shuffle=True)

    # 从无标记样本池中移除选中的样本
    unlabeled_indices = [i for i in unlabeled_indices if i not in selected_unlabeled_indices]
    unlabeled_pool = Subset(train_data, unlabeled_indices)

end_time = time()

print("time consumption: ",end_time - start_time )