在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模型的完整流程。希望对你有所帮助!