超大超详细垃圾分类数据集(分类,分类),共4大类,345小类,147913张图,已全部分类标注完成,共12GB。
厨余垃圾 76小类 35058张
可回收物 195类 86116张
其他垃圾 53类 16156张
有害垃圾 18小类 10583张
如何使用深度学习框架(如PyTorch)来训练一个包含147913张图像的超大超详细垃圾分类数据集,并附上详细的训练代码和步骤。这个数据集包含4大类,345小类,已全部分类标注完成,总大小为12GB。
数据集描述
- 数据量:147913张图像
- 类别:
- 厨余垃圾:76小类,35058张
- 可回收物:195小类,86116张
- 其他垃圾:53小类,16156张
- 有害垃圾:18小类,10583张
- 总大小:12GB
- 任务类型:图像分类
数据集组织
假设你的数据集目录结构如下:
garbage_classification_dataset/
├── train/
│ ├── food_waste/
│ │ ├── class1/
│ │ ├── class2/
│ │ └── ...
│ ├── recyclable/
│ │ ├── class1/
│ │ ├── class2/
│ │ └── ...
│ ├── other_waste/
│ │ ├── class1/
│ │ ├── class2/
│ │ └── ...
│ └── hazardous_waste/
│ ├── class1/
│ ├── class2/
│ └── ...
├── valid/
│ ├── food_waste/
│ │ ├── class1/
│ │ ├── class2/
│ │ └── ...
│ ├── recyclable/
│ │ ├── class1/
│ │ ├── class2/
│ │ └── ...
│ ├── other_waste/
│ │ ├── class1/
│ │ ├── class2/
│ │ └── ...
│ └── hazardous_waste/
│ ├── class1/
│ ├── class2/
│ └── ...
└── test/
├── food_waste/
│ ├── class1/
│ ├── class2/
│ └── ...
├── recyclable/
│ ├── class1/
│ ├── class2/
│ └── ...
├── other_waste/
│ ├── class1/
│ ├── class2/
│ └── ...
└── hazardous_waste/
├── class1/
├── class2/
└── ...
安装依赖
确保你已经安装了必要的依赖库:
pip install torch torchvision matplotlib
数据加载和预处理
使用torchvision
中的ImageFolder
来加载数据集,并进行预处理:
import torch
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
# 定义数据预处理
transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
# 加载数据集
train_dataset = datasets.ImageFolder(root='./garbage_classification_dataset/train', transform=transform)
valid_dataset = datasets.ImageFolder(root='./garbage_classification_dataset/valid', transform=transform)
test_dataset = datasets.ImageFolder(root='./garbage_classification_dataset/test', transform=transform)
# 创建数据加载器
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=4)
valid_loader = DataLoader(valid_dataset, batch_size=32, shuffle=False, num_workers=4)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False, num_workers=4)
模型定义
使用预训练的ResNet50模型进行迁移学习:
import torch.nn as nn
import torchvision.models as models
# 加载预训练的ResNet50模型
model = models.resnet50(pretrained=True)
# 替换最后一层全连接层
num_features = model.fc.in_features
model.fc = nn.Linear(num_features, 345) # 345个类别
# 将模型移动到GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)
损失函数和优化器
定义损失函数和优化器:
import torch.optim as optim
# 定义损失函数
criterion = nn.CrossEntropyLoss()
# 定义优化器
optimizer = optim.Adam(model.parameters(), lr=0.001)
训练模型
编写训练循环:
def train_model(model, train_loader, valid_loader, criterion, optimizer, num_epochs=100):
best_val_accuracy = 0.0
for epoch in range(num_epochs):
# 训练阶段
model.train()
running_loss = 0.0
running_corrects = 0
for inputs, labels in train_loader:
inputs, labels = inputs.to(device), labels.to(device)
# 前向传播
outputs = model(inputs)
loss = criterion(outputs, labels)
# 反向传播和优化
optimizer.zero_grad()
loss.backward()
optimizer.step()
_, preds = torch.max(outputs, 1)
running_loss += loss.item() * inputs.size(0)
running_corrects += torch.sum(preds == labels.data)
train_loss = running_loss / len(train_loader.dataset)
train_accuracy = running_corrects.double() / len(train_loader.dataset)
# 验证阶段
model.eval()
running_loss = 0.0
running_corrects = 0
with torch.no_grad():
for inputs, labels in valid_loader:
inputs, labels = inputs.to(device), labels.to(device)
outputs = model(inputs)
loss = criterion(outputs, labels)
_, preds = torch.max(outputs, 1)
running_loss += loss.item() * inputs.size(0)
running_corrects += torch.sum(preds == labels.data)
val_loss = running_loss / len(valid_loader.dataset)
val_accuracy = running_corrects.double() / len(valid_loader.dataset)
print(f'Epoch {epoch + 1}/{num_epochs}')
print(f'Train Loss: {train_loss:.4f}, Train Acc: {train_accuracy:.4f}')
print(f'Val Loss: {val_loss:.4f}, Val Acc: {val_accuracy:.4f}')
# 保存最佳模型
if val_accuracy > best_val_accuracy:
best_val_accuracy = val_accuracy
torch.save(model.state_dict(), 'best_model.pth')
# 开始训练
train_model(model, train_loader, valid_loader, criterion, optimizer, num_epochs=100)
模型评估
训练完成后,可以使用以下代码评估模型在测试集上的表现:
def evaluate_model(model, test_loader, criterion):
model.load_state_dict(torch.load('best_model.pth'))
model.eval()
running_loss = 0.0
running_corrects = 0
with torch.no_grad():
for inputs, labels in test_loader:
inputs, labels = inputs.to(device), labels.to(device)
outputs = model(inputs)
loss = criterion(outputs, labels)
_, preds = torch.max(outputs, 1)
running_loss += loss.item() * inputs.size(0)
running_corrects += torch.sum(preds == labels.data)
test_loss = running_loss / len(test_loader.dataset)
test_accuracy = running_corrects.double() / len(test_loader.dataset)
print(f'Test Loss: {test_loss:.4f}, Test Acc: {test_accuracy:.4f}')
# 评估模型
evaluate_model(model, test_loader, criterion)
模型预测
你可以使用训练好的模型对新图像进行预测:
import matplotlib.pyplot as plt
def predict_image(image_path, model, transform, class_names):
model.load_state_dict(torch.load('best_model.pth'))
model.eval()
image = Image.open(image_path).convert('RGB')
image_tensor = transform(image).unsqueeze(0).to(device)
with torch.no_grad():
outputs = model(image_tensor)
_, predicted = torch.max(outputs, 1)
predicted_class = class_names[predicted.item()]
plt.imshow(image)
plt.title(f'Predicted: {predicted_class}')
plt.show()
# 获取类别名称
class_names = [item[0].split('/')[-1] for item in train_loader.dataset.class_to_idx.items()]
# 预测新图像
predict_image('path/to/your/image.jpg', model, transform, class_names)
注意事项
- 数据集质量:确保数据集的质量,包括清晰度、标注准确性等。
- 模型选择:可以选择更强大的模型版本(如ResNet101、EfficientNet等)以提高性能。
- 超参数调整:根据实际情况调整超参数,如批量大小(
batch-size
)、学习率(lr
)等。 - 监控性能:训练过程中监控损失函数和准确率指标,确保模型收敛。
通过上述步骤,你可以使用PyTorch来训练一个超大超详细的垃圾分类数据集,并使用训练好的模型进行预测。
标签:loss,训练,torch,loader,147913,running,train,详细,model From: https://blog.csdn.net/2401_88440984/article/details/143421880