基于Pytorch深度学习的脑肿瘤分类识别

 import numpy as np

    
 from pathlib import Path
    
 import random
    
 import torch
    
 import torchvision
    
 from torchvision import transforms
    
 import tqdm
    
 data_dir = Path("./brain_tumor_dataset")
    
 list(data_dir.iterdir())
    
 # 计算该图像数据集的均值和STD
    
 def calculate_dataset_mean(dataloader): 
    
     images, labels = next(iter(dataloader))
    
     return images.mean([0,2,3])
    
 def calculate_dataset_std(dataloader):
    
     images, labels = next(iter(dataloader))
    
     return images.std([0,2,3])
    
  
    
 raw_dataset_transforms = transforms.Compose([
    
     transforms.Resize(255), 
    
     transforms.CenterCrop(225), 
    
     transforms.ToTensor()
    
     ])
    
 raw_dataset = torchvision.datasets.ImageFolder(root = str(data_dir), transform=raw_dataset_transforms)
    
 raw_dataloader = torch.utils.data.DataLoader(raw_dataset, batch_size=len(raw_dataset))
    
 print(f"mean = {calculate_dataset_mean(raw_dataloader)} and std = {calculate_dataset_std(raw_dataloader)}")
    
 '''
    
 创建数据加载器
    
 为训练和验证数据集组合转换
    
 创建应用了适当转换的数据集
    
 使用数据集创建数据加载程序
    
 '''
    
 # 数据集配置
    
 CLASSES = ["no","yes"]
    
 NUMBER_OF_CLASSES = len(CLASSES)
    
 SHUFFLE = True
    
 VALIDATION_SIZE = 0.2
    
 RESIZE = 64
    
 # 为训练和验证数据集组合图像转换
    
 normalize = transforms.Normalize(
    
     mean=calculate_dataset_mean(raw_dataloader),
    
     std=calculate_dataset_std(raw_dataloader),
    
     )
    
 training_transform = transforms.Compose([
    
     transforms.RandomRotation(30),
    
     transforms.Resize(RESIZE),
    
     transforms.CenterCrop(RESIZE),
    
     transforms.RandomHorizontalFlip(),
    
     transforms.ToTensor(),
    
     normalize,
    
 ] )
    
 validation_transform = transforms.Compose(
    
     [transforms.Resize(RESIZE),
    
      transforms.CenterCrop(RESIZE),         
    
      transforms.ToTensor(),
    
      normalize,
    
     ])
    
  
    
 # 创建数据集并应用相关转换
    
 training_dataset = torchvision.datasets.ImageFolder(
    
     root = str(data_dir),
    
     transform = training_transform,
    
 )
    
 validation_dataset = torchvision.datasets.ImageFolder(
    
     root = str(data_dir),
    
     transform = validation_transform
    
 )
    
 # 检查数据集规范化
    
 # 归一化后，均值应接近0,std应接近1
    
 normalized_dataloader = torch.utils.data.DataLoader(
    
     validation_dataset, 
    
     batch_size=len(validation_dataset),
    
     )
    
 print(f"mean = {calculate_dataset_mean(normalized_dataloader)}",
    
       f"std = {calculate_dataset_std(normalized_dataloader)}")
    
 # 使用采样器在训练和验证之间分割数据
    
 split = int(np.floor(len(training_dataset) * VALIDATION_SIZE))
    
 indices = list(range(len(training_dataset)))
    
 if SHUFFLE:
    
     random.shuffle(indices)
    
  
    
 validation_indices, training_indices = indices[: split], indices[split :]
    
  
    
 training_sampler = torch.utils.data.sampler.SubsetRandomSampler(training_indices)
    
 validation_sampler = torch.utils.data.sampler.SubsetRandomSampler(validation_indices)
    
 # 数据加载器配置
    
 BATCH_SIZE = 4
    
 NUMBER_OF_WORKERS = 2
    
 PIN_MEMORY = False
    
 # 创建数据加载器
    
 training_dataloader = torch.utils.data.DataLoader(
    
     training_dataset, 
    
     batch_size=BATCH_SIZE, 
    
     sampler=training_sampler,
    
     num_workers=NUMBER_OF_WORKERS, 
    
     pin_memory=PIN_MEMORY,
    
     )
    
  
    
 validation_dataloader = torch.utils.data.DataLoader(
    
     training_dataset, 
    
     batch_size=BATCH_SIZE, 
    
     sampler=validation_sampler,
    
     num_workers=NUMBER_OF_WORKERS, 
    
     pin_memory=PIN_MEMORY,
    
     )
    
 import matplotlib.pyplot as plt
    
 %matplotlib inline
    
  
    
 # 展示图片
    
 for images, labels in training_dataloader:
    
     fig = plt.figure(figsize = (14, 7))
    
     for i in range(BATCH_SIZE):
    
     ax = fig.add_subplot(2, 4, i + 1, xticks = [], yticks = [])
    
     ax.set_xlabel(f"cancer = {CLASSES[labels[i]]}")
    
     image = images[i][0, :, :]
    
     plt.imshow(image)
    
     break
    
 # 训练模型配置
    
 MODEL_NAME = "resnet18"
    
 WEIGHTS = "DEFAULT"
    
 LEARNING_RATE = 0.0001
    
 MOMENTUM = 0.9
    
 NUMBER_OF_EPOCHS = 10
    
 MODEL_SAVE_PATH = "model.pt"
    
  
    
 # 初始化模型
    
 model = torchvision.models.get_model(MODEL_NAME, weights=WEIGHTS)
    
 model.fc = torch.nn.Linear(512, 2)
    
 # 选择一个损失函数和一个优化函数
    
 criterion = torch.nn.CrossEntropyLoss()
    
 optimizer = torch.optim.SGD(model.parameters(), lr=LEARNING_RATE, momentum=MOMENTUM)
    
 # 开始训练模型
    
 # 存储训练过程状态的地方
    
 training_loss_history = []
    
 validation_loss_history = []
    
 validation_accuracy_history = []
    
  
    
 # 训练和验证循环
    
 for epoch in range(NUMBER_OF_EPOCHS):
    
     
    
     model.train()
    
     training_step_loss = []
    
     print(f"Epoch {epoch + 1}/{NUMBER_OF_EPOCHS}")
    
     for data in tqdm.tqdm(training_dataloader, desc="training"):
    
     features, labels = data
    
     optimizer.zero_grad()
    
     outputs = model(features)
    
     training_loss = criterion(outputs, labels)
    
     training_loss.backward()
    
     optimizer.step()
    
     training_step_loss.append(training_loss.item())
    
         
    
     training_epoch_loss = sum(training_step_loss)/len(training_step_loss)
    
     training_loss_history.append(training_epoch_loss)
    
     
    
     model.eval()
    
     validation_step_loss = []
    
     correct_predictions = 0 
    
     for data in tqdm.tqdm(validation_dataloader, desc="validating"):
    
  
    
     features, labels = data
    
     outputs = model(features)
    
     correct_predictions += torch.sum(torch.argmax(outputs, axis=1)==labels)
    
     validation_loss = criterion(outputs, labels)
    
     validation_step_loss.append(validation_loss.item())
    
  
    
     validation_epoch_loss = sum(validation_step_loss)/len(validation_step_loss)
    
     validation_loss_history.append(validation_epoch_loss)
    
     validation_epoch_accuracy = correct_predictions / (len(validation_dataloader) * BATCH_SIZE)
    
     print(f"Training Loss: {training_epoch_loss:.4f},"
    
       f"Validation Loss: {validation_epoch_loss:.4f}," 
    
       f"Validation Acc: {validation_epoch_accuracy:.4f}")
    
     
    
     # 保存模型
    
     if epoch==0 or validation_epoch_accuracy > max(validation_accuracy_history):
    
     print("Validation loss improved, saving checkpoint.")
    
     torch.save({
    
             'epoch': epoch,
    
             'model_state_dict': model.state_dict(),
    
             'optimizer_state_dict': optimizer.state_dict(),
    
             'loss': validation_epoch_loss,
    
             }, MODEL_SAVE_PATH)
    
     print("Checkpoint saved")
    
     validation_accuracy_history.append(validation_epoch_accuracy)
    
 print('Finished Training')
    
 # 绘制训练和验证损失
    
 plt.plot(training_loss_history, label='training_loss')
    
 plt.plot(validation_loss_history,label='validation_loss')
    
 plt.legend()
    
 plt.show
    
 # 绘制验证精度图
    
 plt.plot(validation_accuracy_history, label='validation accuracy')
    
 plt.legend()
    
 plt.show