如何利用RseNet训练自己的图像识别模型?

356 阅读2分钟

在ResNet基础上增加自己的训练集并进行微调,然后保存和调用微调后的模型,可以按照以下步骤操作:

1. 准备数据集

  • 组织数据:确保你的数据集按照类别组织好,每个类别的图片放在一个单独的文件夹中。
  • 数据增强:可以使用数据增强技术来提高模型的泛化能力。

2. 加载预训练的ResNet模型

import torch
import torchvision.models as models

# 加载预训练的ResNet模型
model = models.resnet50(pretrained=True)

3. 修改模型以适应自己的任务

  • 修改分类层:根据你的数据集类别数修改最后一层全连接层。
num_classes = 10  # 假设你的数据集有10个类别
model.fc = torch.nn.Linear(model.fc.in_features, num_classes)

4. 准备数据加载器

from torchvision import datasets, transforms
from torch.utils.data import DataLoader

# 数据变换
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# 加载数据集
train_dataset = datasets.ImageFolder('path/to/train', transform=transform)
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)

val_dataset = datasets.ImageFolder('path/to/val', transform=transform)
val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False)

5. 训练模型

import torch.optim as optim
import torch.nn as nn

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

# 训练模型
num_epochs = 10
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

for epoch in range(num_epochs):
    model.train()
    running_loss = 0.0
    for inputs, labels in train_loader:
        inputs, labels = inputs.to(device), labels.to(device)
        
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        
        running_loss += loss.item()
    
    print(f'Epoch {epoch+1}/{num_epochs}, Loss: {running_loss/len(train_loader)}')

6. 保存模型

torch.save(model.state_dict(), 'resnet50_finetuned.pth')

7. 调用微调后的模型

# 加载微调后的模型
model = models.resnet50(pretrained=False)
model.fc = torch.nn.Linear(model.fc.in_features, num_classes)
model.load_state_dict(torch.load('resnet50_finetuned.pth'))
model.to(device)
model.eval()

8. 使用模型进行预测

from PIL import Image

def predict(image_path):
    image = Image.open(image_path).convert('RGB')
    image = transform(image).unsqueeze(0).to(device)
    with torch.no_grad():
        output = model(image)
        _, predicted = torch.max(output, 1)
    return predicted.item()

# 示例
image_path = 'path/to/your/image.jpg'
predicted_class = predict(image_path)
print(f'Predicted class: {predicted_class}')

以上步骤涵盖了从准备数据集到保存和调用微调后的ResNet模型的完整流程。希望对你有所帮助!