Datawhale-AI夏令营:脑PET图像分析和疾病预测挑战赛

import os, sys, glob, argparse
import pandas as pd
import numpy as np
from tqdm import tqdm

import cv2
from PIL import Image
from sklearn.model_selection import train_test_split, StratifiedKFold, KFold

import torch
torch.manual_seed(0)
torch.backends.cudnn.deterministic = False
torch.backends.cudnn.benchmark = True

import torchvision.models as models
import torchvision.transforms as transforms
import torchvision.datasets as datasets
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.autograd import Variable
from torch.utils.data.dataset import Dataset

import nibabel as nib
from nibabel.viewers import OrthoSlicer3D

#数据集路径
train_path = glob.glob('./PET/PET/Train/*/*')#返回所有匹配的文件列表,这里是图片路径列表
test_path = glob.glob('./PET/PET/Test/*')

#随机打乱图片路径
np.random.shuffle(train_path)
np.random.shuffle(test_path)

DATA_CACHE = {}
#加载数据
class XunFeiDataset(Dataset):
    def __init__(self, img_path, transform=None):
        """初始化图片路径和数据增强操作"""
        self.img_path = img_path
        #使用transform做数据增强
        if transform is not None:
            self.transform = transform
        else:
            self.transform = None
    
    def __getitem__(self, index):
        """实现对象的索引操作"""
        #若图片在DATA_CACHE中，就使用DATA_CACHE中的
        #否则加载到DATA_CACHE中
        if self.img_path[index] in DATA_CACHE:
            img = DATA_CACHE[self.img_path[index]]
        else:
            img = nib.load(self.img_path[index]) 
            img = img.dataobj[:,:,:, 0]
            DATA_CACHE[self.img_path[index]] = img
        
        # 随机选择一些通道            
        idx = np.random.choice(range(img.shape[-1]), 50)
        img = img[:, :, idx]
        img = img.astype(np.float32)

        if self.transform is not None:
            img = self.transform(image = img)['image']
        
        #transpose将数据的格式重新排列
        #原始格式img[H,W,C],H:hight图片高度;W:weight图片宽度;C:channel图片通道数
        #transpose后变为img[C,H,W]
        #为什么要这样换？
        #因为神经网络要求的输入图片的格式是这样的
        img = img.transpose([2,0,1])
        #将numpy_array转换成tensor，因为神经网络处理的是tensor
        #print(img.shape) [50,120,120]
        return img,torch.from_numpy(np.array(int('NC' in self.img_path[index])))
    
    def __len__(self):
        """获取数据集图片数量"""
        return len(self.img_path)
#albumentations是一个数据增强库  
#Albumentations相比于torchvision.transforms
# 提供了更多的图像增强方法和更高的灵活性
# 具体的操作效果可以自行去csdn搜索      
import albumentations as A
#划分训练集和验证集
#训练集是从头到倒数第十个路径之前的图片
train_loader = torch.utils.data.DataLoader(
    #截取 train_path 中从头到倒数第十个路径之前的图片路径
    XunFeiDataset(train_path[:-10],
            #将多个数据增强操作组合在一起，作为实参传给transform这个形参
            A.Compose([ 
            #随机旋转90°
            A.RandomRotate90(),
            # 随机裁剪
            A.RandomCrop(120, 120),
            #围绕Y轴水平翻转
            A.HorizontalFlip(p=0.5),
            A.RandomContrast(p=0.5),
            #随机亮度对比度
            A.RandomBrightnessContrast(p=0.5),
        ])
    ), batch_size=2, shuffle=True, num_workers=1, pin_memory=False
)
#batch_size:批大小，指一次送入网路中训练的图片数量,视电脑显存而定
#num_workers:使用几个进程读取数据
#pin_memory:指是否将加载的数据存储在 CUDA 固定内存中
#验证集是最后10张图片
val_loader = torch.utils.data.DataLoader(
    XunFeiDataset(train_path[-10:],
            A.Compose([
            A.RandomCrop(120, 120),
        ])
    ), batch_size=2, shuffle=False, num_workers=1, pin_memory=False
)
#加载测试集
test_loader = torch.utils.data.DataLoader(
    XunFeiDataset(test_path,
            A.Compose([
            A.RandomCrop(128, 128),
            A.HorizontalFlip(p=0.5),
            A.RandomContrast(p=0.5),
        ])
    ), batch_size=2, shuffle=False, num_workers=1, pin_memory=False
)
"""
补充：
    1.训练集、验证集、测试集：训练神经网络时要划分训练集、验证集和测试集。
    训练集用来训练模型，测试集用来最终测试模型的好坏，例如高考。验证集用来测试每次训练的模型的效果，例如在学校的平时测试。
    2.划分方法：通常训练集与测试集进行8、2分，即训练集占总数据的80%，测试集占20%，在训练集中再划分验证集
    3.K折交叉验证：若在训练集中划分K次训练集和验证集，每次都将训练集分为K份，每次训练时选择K-1份作为训练集，1份作为验证集，
    称为K折交叉验证。训练后取K个训练模型的平均结果作为最终结果，可能会提高模型效果。
"""

#定义模型架构
class XunFeiNet(nn.Module):
    def __init__(self):
        super(XunFeiNet, self).__init__()
        #使用预训练的resnet18模型        
        model = models.resnet18(True)
        #定义了一个卷积层，in_channel=50，out_channel=64,卷积核大小7X7,步长为2，填充为6(左右各3)
        model.conv1 = torch.nn.Conv2d(50, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
        #adaptiveavgpool2d(1)是一个自适应平均池化层，
        # 它可以根据输入的大小自动调整池化核的大小，将输入的特征图压缩成一个大小为1x1的特征图。
        model.avgpool = nn.AdaptiveAvgPool2d(1)
        #为什么线性层的输入是512，最后的线性层输出为什么是2？
        #因为进入线性层前需要将图片展开拉平，512=C*H*W,输出是2因为做的是二分类
        model.fc = nn.Linear(512, 2)
        self.resnet = model
    #前向传播，这里的反向传播和梯度更新pytorch都封装好了，想要了解的同学可以看下面的书    
    def forward(self, img):        
        out = self.resnet(img)
        return out
        
model = XunFeiNet()
#使用GPU训练模型
model = model.to('cuda')
#使用交叉熵损失函数
criterion = nn.CrossEntropyLoss().cuda()
#定义参数更新部分的优化迭代器，学习率设为0.001
optimizer = torch.optim.AdamW(model.parameters(), 0.001)
"""
补充几个需要了解的概念：
1.CNN都有哪些层？每个层都有什么作用？如何用pytorch定义？
2.激活函数有哪些？现在常用的有哪些？resnet为什么用的是ReLU(上述代码是直接调用的resnet所以没有显示)？
3.resnet的具体结构是怎样的，resnet18,resnet50,resnet101的区别？跳跃连接的作用是什么？为什么要用1x1的卷积去实现跳跃连接？
4.参数更新的优化方法有哪些？
5.学习率(lr)、批大小(batch_size)、填充(padding)、步长(stride)、每一层神经元的数量(out_channel)这些超参数有什么意义？
尝试去修改这些超参数。给出一个输入图片，你能算出通过卷积层和池化层后的图片大小吗？
上述内容不了解的同学可以去看《动手学深度学习》链接如下：
https://zh-v2.d2l.ai/chapter_convolutional-neural-networks/index.html
"""

#训练模型
def train(train_loader, model, criterion, optimizer):
    model.train()
    train_loss = 0.0
    for i, (input, target) in enumerate(train_loader):
        input = input.cuda(non_blocking=True)
        target = target.cuda(non_blocking=True)
        #调用model进行前向传播，这里没有调用forward()函数是怎么实现前向传播的呢？
        #因为我们的网络继承了nn.Module这个父类，module的call里面调用module的forward方法
        #可以了解一下__call__函数的作用
        output = model(input)
        loss = criterion(output, target)
        #optimizer.zero_grad()将模型中所有可训练的参数的梯度清零。
        # 在训练神经网络时，通常需要在每次迭代之前调用这个函数。因为如果不清零梯度，
        # 优化器在更新权重时会累加之前的梯度，导致梯度计算错误，不能正确更新参数
        optimizer.zero_grad()
        #反向传播
        loss.backward()
        #使用前面定义的AdamW方法更新参数如：权重(weight)、偏差(bias)
        optimizer.step()
        #每20个数据打印一遍loss
        if i % 20 == 0:
            print(loss.item())
        #累计一个batch的loss    
        train_loss += loss.item()
    #返回一个batch的平均loss
    return train_loss/len(train_loader)

#验证模型效果           
def validate(val_loader, model, criterion):
    model.eval()
    val_acc = 0.0
    #with torch.no_grad()表示下面所有操作都不会被追踪以用于求导。
    # 因为验证和预测时网络的参数已经固定好了，只需要前向传播，不需要反向传播和参数更新
    #同时这样可以节省内存和加速计算
    with torch.no_grad():
        #enumerate将一个 dataloader 对象作为输入，并返回一个可迭代对象，
        # 该对象产生一对 (index, data) 元组，
        # 其中 index 是当前 batch 的索引，data 是当前 batch 的数据。
        for i, (input, target) in enumerate(val_loader):
            input = input.cuda()
            target = target.cuda()

            # compute output
            output = model(input)
            loss = criterion(output, target)
            
            val_acc += (output.argmax(1) == target).sum().item()
            
    return val_acc / len(val_loader.dataset)

#训练3个epoch(轮)
# 一般模型的准确率会随着epoch的增加而增加，
# 但训练次数过多有可能会导致过拟合
# 少了又可能会欠拟合    
for _  in range(3):
    train_loss = train(train_loader, model, criterion, optimizer)
    val_acc  = validate(val_loader, model, criterion)
    train_acc = validate(train_loader, model, criterion)
    
    print(train_loss, train_acc, val_acc)
"""
补充：
1.什么是过拟合和欠拟合？如何区别？
2.过拟合和欠拟合都有哪些解决方法？
3.训练模型的流程都有一个固定的套路
"""

#测试
#训练、验证、测试的代码看起来大同小异，大家一起来找茬
def predict(test_loader, model, criterion):
    model.eval()
    val_acc = 0.0
    
    test_pred = []
    with torch.no_grad():
        for i, (input, target) in enumerate(test_loader):
            input = input.cuda()
            target = target.cuda()

            output = model(input)
            test_pred.append(output.data.cpu().numpy())
            
    return np.vstack(test_pred)
    
pred = None
for _ in range(10):
    if pred is None:
        pred = predict(test_loader, model, criterion)
    else:
        pred += predict(test_loader, model, criterion)
        
submit = pd.DataFrame(
    {
        #测试数据的标识符（通过解析测试数据路径中的文件名得到）
        'uuid': [int(x.split('/')[-1][:-4]) for x in test_path],
        #返回每个样本的最大值所在的索引（即预测标签的索引）。
        'label': pred.argmax(1)
})
submit['label'] = submit['label'].map({1:'NC', 0: 'MCI'})
submit = submit.sort_values(by='uuid')
submit.to_csv('submit2.csv', index=None)