PyTorch下的5种不同神经网络-ResNet

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((224, 224)),  # ResNet的输入图像大小

    transforms.ToTensor(),

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

])

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.加载预训练的ResNet模型

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

resnet_model = models.resnet18(pretrained=True)

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

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

num_ftrs = resnet_model.fc.in_features

resnet_model.fc = nn.Linear(num_ftrs, len(dataset.label_to_index))

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

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

criterion = nn.CrossEntropyLoss()

optimizer = optim.Adam(resnet_model.parameters(), lr=0.001)

9.模型并行化

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

if torch.cuda.device_count() > 1:

    resnet_model = nn.DataParallel(resnet_model)

10.将模型发送到GPU

模型发送到GPU进行训练

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

resnet_model.to(device)

11.训练模型

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

num_epochs = 10

for epoch in range(num_epochs):

    resnet_model.train()

    running_loss = 0.0

    for images, labels in data_loader:

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



        # 前向传播

        outputs = resnet_model(images)

        loss = criterion(outputs, labels)



        # 反向传播和优化

        optimizer.zero_grad()

        loss.backward()

        optimizer.step()



        running_loss += loss.item()



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