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}')

  • 8
    点赞
  • 6
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值