PyTorch学习笔记：新冠肺炎X光分类

    ├─imgs
    │  ├─covid-19
    │  │      auntminnie-a-2020_01_28_23_51_6665_2020_01_28_Vietnam_coronavirus.jpeg
    │  │      ryct.2020200028.fig1a.jpeg
    │  │
    │  └─no-finding
    │         00001215_000.png
    │         00001215_001.png
    │
    └─labels
       train.txt
       valid.txt
    
    
      
      
      
      
      
      
      
      
      
      
      
      
    

    covid-19/auntminnie-a-2020_01_28_23_51_6665_2020_01_28_Vietnam_coronavirus.jpeg 1
    no-finding/00001215_000.png 0
    
    
    
      
      
      
    

    covid-19/ryct.2020200028.fig1a.jpeg 1
    no-finding/00001215_001.png 0
    
    
      
      
    

    import os
    from PIL import Image
    from torch.utils.data import Dataset, DataLoader
    from torchvision import transforms
    import torch
    import torch.nn as nn
    import torch.optim as optim
    import torch.nn.functional as F
    
    # 自定义数据集类
    class COVID19Dataset(Dataset):
    def __init__(self, img_dir, label_file, transform=None):
        self.img_dir = img_dir
        self.transform = transform
        self.img_labels = []
    
        with open(label_file, 'r') as f:
            lines = f.readlines()
            for line in lines:
                self.img_labels.append(line.strip().split())
    
    def __len__(self):
        return len(self.img_labels)
    
    def __getitem__(self, idx):
        img_path, label = self.img_labels[idx]
        img_path = os.path.join(self.img_dir, img_path)
        image = Image.open(img_path).convert('RGB')
        label = int(label)
    
        if self.transform:
            image = self.transform(image)
    
        return image, label
    
    # 图像预处理
    transform = transforms.Compose([
    transforms.Resize((8, 8)),
    transforms.ToTensor()
    ])
    
    # 创建数据集和数据加载器
    train_dataset = COVID19Dataset(img_dir='imgs', label_file='labels/train.txt', transform=transform)
    train_loader = DataLoader(train_dataset, batch_size=2, shuffle=True)
    
    valid_dataset = COVID19Dataset(img_dir='imgs', label_file='labels/valid.txt', transform=transform)
    valid_loader = DataLoader(valid_dataset, batch_size=2, shuffle=False)
    
    # 定义简单卷积神经网络
    class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        self.conv1 = nn.Conv2d(3, 1, kernel_size=3)  # 输入通道为3（RGB），输出通道为1，卷积核大小为3x3
        self.fc1 = nn.Linear(1 * 6 * 6, 2)  # 全连接层，输入大小为6*6*1，输出大小为2（2类）
    
    def forward(self, x):
        x = self.conv1(x)
        x = F.relu(x)
        x = x.view(-1, 1 * 6 * 6)  # 展平操作
        x = self.fc1(x)
        return x
    
    model = SimpleCNN()
    
    # 定义损失函数和优化器
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(model.parameters(), lr=0.01)
    
    # 训练函数
    def train(model, train_loader, criterion, optimizer, epoch):
    model.train()
    running_loss = 0.0
    for batch_idx, (data, target) in enumerate(train_loader):
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
        if batch_idx % 10 == 9:
            print(f'Train Epoch: {epoch} [{batch_idx * len(data)}/{len(train_loader.dataset)}] Loss: {running_loss / 10:.6f}')
            running_loss = 0.0
    
    # 验证函数
    def validate(model, valid_loader, criterion):
    model.eval()
    validation_loss = 0.0
    correct = 0
    with torch.no_grad():
        for data, target in valid_loader:
            output = model(data)
            validation_loss += criterion(output, target).item()
            pred = output.argmax(dim=1, keepdim=True)
            correct += pred.eq(target.view_as(pred)).sum().item()
    
    validation_loss /= len(valid_loader.dataset)
    print(f'\nValidation set: Average loss: {validation_loss:.4f}, Accuracy: {correct}/{len(valid_loader.dataset)} ({100. * correct / len(valid_loader.dataset):.0f}%)\n')
    
    # 训练和验证
    for epoch in range(1, 11):
    train(model, train_loader, criterion, optimizer, epoch)
    validate(model, valid_loader, criterion)