PyTorch下的5种不同神经网络-一.AlexNet

1.导入模块

导入所需的Python库，包括图像处理、深度学习模型和数据加载

import os

import torch

import torch.nn as nn

import torch.optim as optim

from torch.utils.data import Dataset, DataLoader

from PIL import Image

from torchvision import models, transforms

2.定义自定义图像数据集

创建一个自定义的图像数据集类，用于加载和处理图像数据。

class CustomImageDataset(Dataset):

    def __init__(self, main_dir, transform=None):

        self.main_dir = main_dir

        self.transform = transform

        self.files = []

        self.labels = []

        self.label_to_index = {}

        for index, label in enumerate(os.listdir(main_dir)):

            self.label_to_index[label] = index

            label_dir = os.path.join(main_dir, label)

            if os.path.isdir(label_dir):

                for file in os.listdir(label_dir):

                    self.files.append(os.path.join(label_dir, file))

                    self.labels.append(label)



    def __len__(self):

        return len(self.files)



    def __getitem__(self, idx):

        image = Image.open(self.files[idx])

        label = self.labels[idx]

        if self.transform:

            image = self.transform(image)

        return image, self.label_to_index[label]

3.定义数据转换

定义一个数据转换过程，包括图像大小调整、随机翻转、旋转、转换为张量以及标准化

transform = transforms.Compose([

    transforms.Resize((227, 227)),  # AlexNet的输入图像大小

    transforms.RandomHorizontalFlip(),  # 随机水平翻转

    transforms.RandomRotation(10),  # 随机旋转

    transforms.ToTensor(),

    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),  # AlexNet的标准化

])

4.创建数据集

使用自定义数据集类和定义的数据转换来创建数据集

dataset = CustomImageDataset(main_dir="F:\\A-GX\\A-SJJ\\flower_photos\\flower_photos", transform=transform)

5.创建数据加载器

使用数据集创建一个数据加载器，用于批量加载和处理数据

data_loader = DataLoader(dataset, batch_size=32, shuffle=True)

6.加载预训练的AlexNet模型

从PyTorch库中加载预训练的AlexNet模型

alexnet_model = models.alexnet(pretrained=True)

7.修改最后几层以适应新的分类任务

修改AlexNet模型的最后几层，以便它能够处理新的分类任务

num_ftrs = alexnet_model.classifier[6].in_features

alexnet_model.classifier[6] = nn.Linear(num_ftrs, len(dataset.label_to_index))

8.定义损失函数和优化器

定义用于训练模型的损失函数和优化器。

criterion = nn.CrossEntropyLoss()

optimizer = optim.Adam(alexnet_model.parameters(), lr=0.0001)

9.模型并行化

如果有多GPU，则使用nn.DataParallel来并行化模型

if torch.cuda.device_count() > 1:

    alexnet_model = nn.DataParallel(alexnet_model)

10.将模型发送到GPU

将模型发送到GPU进行训练。

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

alexnet_model.to(device)

11.训练模型

数据加载器和定义的参数训练模型

num_epochs = 10

for epoch in range(num_epochs):

    alexnet_model.train()

    running_loss = 0.0

    for images, labels in data_loader:

        images, labels = images.to(device), labels.to(device)



        # 前向传播

        outputs = alexnet_model(images)

        loss = criterion(outputs, labels)



        # 反向传播和优化

        optimizer.zero_grad()

        loss.backward()

        optimizer.step()



        running_loss += loss.item()



    # 在每个epoch结束后评估模型

    train_accuracy = evaluate_model(alexnet_model, data_loader, device)

    print(f'Epoch [{epoch + 1}/{num_epochs}], Loss: {running_loss / len(data_loader):.4f}, Train Accuracy: {train_accuracy:.2f}%')

12.评估模型

定义一个评估函数，用于评估模型的性能

def evaluate_model(model, data_loader, device):

    model.eval()  # 将模型设置为评估模式

    correct = 0

    total = 0

    with torch.no_grad():  # 在这个块中，所有计算都不会计算梯度

        for images, labels in data_loader:

            images, labels = images.to(device), labels.to(device)

            outputs = model(images)

            _, predicted = torch.max(outputs.data, 1)

            total += labels.size(0)

            correct += (predicted == labels).sum().item()

    accuracy = 100 * correct / total

return accuracy