Compare commits
3 commits
968e108857
...
808191b843
| Author | SHA1 | Date | |
|---|---|---|---|
| 808191b843 | |||
| 6f54c8b13e | |||
| f8bb340a70 |
2 changed files with 209 additions and 36 deletions
153
README.md
153
README.md
|
|
@ -1,14 +1,157 @@
|
|||
# trash-division
|
||||
## 一个基于卷积神经网络的垃圾分类识别系统
|
||||
|
||||
### 同济大学python人工智能程序设计课程小组作业
|
||||
一个基于卷积神经网络的垃圾分类识别系统
|
||||
|
||||
> 同济大学 Python 人工智能程序设计课程小组作业
|
||||
|
||||
基于自定义 ResNet 风格 Bottleneck 架构的 CNN 模型（约 80M 参数），将生活垃圾分为厨余垃圾、可回收物、其他垃圾、有害垃圾四个类别，输入为 256×256 RGB 图像。
|
||||
|
||||
---
|
||||
|
||||
## 目录
|
||||
|
||||
- [项目特点](#项目特点)
|
||||
- [模型架构](#模型架构)
|
||||
- [数据集](#数据集)
|
||||
- [环境要求](#环境要求)
|
||||
- [快速开始](#快速开始)
|
||||
- [文件说明](#文件说明)
|
||||
- [目录结构](#目录结构)
|
||||
- [训练细节](#训练细节)
|
||||
- [许可证](#许可证)
|
||||
|
||||
---
|
||||
|
||||
## 项目特点
|
||||
|
||||
- **四类垃圾分类**：厨余垃圾（1）、可回收物（2）、其他垃圾（3）、有害垃圾（4）
|
||||
- **自定义 ResNet Bottleneck 架构**：约 80M 参数，50 层深度残差网络
|
||||
- **数据增强**：训练时使用随机裁剪、水平翻转、旋转、色彩抖动
|
||||
- **Macro-F1 评估**：采用宏平均 F1 分数作为主要评估指标，兼顾各类别表现
|
||||
- **类别加权损失**：自动计算类别权重，缓解类别不平衡问题
|
||||
- **余弦退火学习率调度**：使用 CosineAnnealingLR 平滑调整学习率
|
||||
- **断点续训**：自动检测 `best_model.pth` 并加载继续训练
|
||||
- **多设备支持**：自动选择 CUDA > Intel XPU > CPU
|
||||
|
||||
## 模型架构
|
||||
|
||||
模型基于残差网络（ResNet）的 Bottleneck 构建块设计。
|
||||
|
||||
### Bottleneck 块
|
||||
|
||||
每个 Bottleneck 块包含三个卷积层：
|
||||
|
||||
| 层 | 卷积 | 作用 |
|
||||
|---|---|---|
|
||||
| 1x1 Conv | 降维 | 减少通道数，降低计算量 |
|
||||
| 3x3 Conv | 特征提取 | 核心卷积操作 |
|
||||
| 1x1 Conv | 升维 (x4) | 恢复通道数至输入的 4 倍 |
|
||||
|
||||
### 网络结构
|
||||
|
||||
| 阶段 | 块数 | 输出通道数 | 说明 |
|
||||
|---|---|---|---|
|
||||
| 初始层 | - | 64 | 7x7 Conv, stride=2 + MaxPool |
|
||||
| Stage 1 | 3 | 256 | 第一个残差阶段 |
|
||||
| Stage 2 | 4 | 512 | - |
|
||||
| Stage 3 | 14 | 1024 | 最深阶段（比 ResNet-50 加深） |
|
||||
| Stage 4 | 3 | 2048 | 最终残差阶段 |
|
||||
| 分类头 | - | 4 | 全局平均池化 + 全连接层 |
|
||||
|
||||
## 数据集
|
||||
|
||||
本项目使用 [tany0699/garbage265](https://modelscope.cn/datasets/tany0699/garbage265) 中文生活垃圾分类数据集，包含 265 个子类别的生活垃圾图片。
|
||||
|
||||
通过 `Merge_classes.py` 脚本将 265 个子类别合并为 4 个顶级类别：
|
||||
|
||||
```
|
||||
厨余垃圾 -> 1
|
||||
可回收物 -> 2
|
||||
其他垃圾 -> 3
|
||||
有害垃圾 -> 4
|
||||
```
|
||||
|
||||
数据集预期放置在 `../trash_division_data/`（与项目根目录平级的兄弟目录）。
|
||||
|
||||
## 环境要求
|
||||
|
||||
本项目无 `requirements.txt`，需手动安装以下依赖：
|
||||
|
||||
- Python 3.8+
|
||||
- PyTorch（推荐 1.10+）
|
||||
- torchvision
|
||||
- tqdm
|
||||
- matplotlib
|
||||
- pandas
|
||||
- Pillow
|
||||
- torchsummary
|
||||
|
||||
## 快速开始
|
||||
|
||||
1. **数据预处理**：将 265 个子类别合并为 4 个顶级类别
|
||||
|
||||
```bash
|
||||
python Merge_classes.py
|
||||
```
|
||||
|
||||
2. **训练模型**：
|
||||
|
||||
```bash
|
||||
python Train.py
|
||||
```
|
||||
|
||||
> **注意**：
|
||||
> - 数据目录默认为 `../trash_division_data/ultimate_4_class/`，需先运行合并脚本
|
||||
> - Windows 系统需将 `num_workers` 设为 `0`（参见 `Dataloader.py` 和 `Train.py`）
|
||||
> - 训练会自动从 `best_model.pth` 断点续训（若存在）
|
||||
|
||||
## 文件说明
|
||||
|
||||
| 文件 | 功能 |
|
||||
|---|---|
|
||||
| `Train.py` | 训练主脚本，包含训练循环、验证、评估 |
|
||||
| `Dataloader.py` | 数据加载模块，包含 RobustImageFolder 和 DataLoader 创建 |
|
||||
| `Model.py` | 模型定义，Bottleneck 残差块 + Net 主模型 |
|
||||
| `Merge_classes.py` | 数据集预处理，265 类合并为 4 类 |
|
||||
| `best_model.pth` | 训练好的最佳模型权重（约 125 MB） |
|
||||
| `AGENTS.md` | AI 助手指南（开发辅助） |
|
||||
| `THIRD_PARTY_LICENSES.md` | 第三方数据集许可证声明 |
|
||||
|
||||
## 目录结构
|
||||
|
||||
```
|
||||
trash-division/
|
||||
├── AGENTS.md # AI 助手指南
|
||||
├── best_model.pth # 最佳模型权重
|
||||
├── Dataloader.py # 数据加载模块
|
||||
├── .gitattributes # Git 属性配置
|
||||
├── LICENSE # MIT 许可证
|
||||
├── Merge_classes.py # 数据集预处理脚本
|
||||
├── Model.py # 模型定义
|
||||
├── README.md # 项目说明（本文件）
|
||||
├── THIRD_PARTY_LICENSES.md # 第三方许可证声明
|
||||
└── Train.py # 训练主脚本
|
||||
```
|
||||
|
||||
## 训练细节
|
||||
|
||||
| 配置项 | 说明 |
|
||||
|---|---|
|
||||
| 输入尺寸 | 256 x 256 RGB |
|
||||
| 优化器 | SGD（momentum=0.9, weight_decay=1e-4） |
|
||||
| 初始学习率 | 0.001 |
|
||||
| 学习率调度 | CosineAnnealingLR |
|
||||
| 损失函数 | 类别加权 CrossEntropyLoss |
|
||||
| 评估指标 | Macro-F1（宏平均 F1 分数） |
|
||||
| 批量大小 | 默认 16（可通过参数调整） |
|
||||
| 训练轮数 | 默认 20（可通过参数调整） |
|
||||
| 设备选择优先级 | CUDA > Intel XPU > CPU |
|
||||
| 断点续训 | 自动检测 best_model.pth 并加载 |
|
||||
|
||||
训练时数据增强管线：RandomResizedCrop(256, scale=(0.8, 1.0)) + RandomHorizontalFlip(p=0.5) + RandomRotation(+-15 deg) + ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2)
|
||||
|
||||
## 许可证
|
||||
|
||||
本项目主代码采用 [MIT 许可证](LICENSE)。
|
||||
|
||||
本项目包含的数据集 `tany0699/garbage265` 采用 [Apache License 2.0](THIRD_PARTY_LICENSES.md)，详情请参阅 `THIRD_PARTY_LICENSES.md` 文件。
|
||||
|
||||
|
||||
|
||||
|
|
|
|||
92
Train.py
92
Train.py
|
|
@ -14,51 +14,67 @@ import matplotlib.pyplot as plt
|
|||
from Model import Net
|
||||
from Dataloader import create_dataloaders
|
||||
import os
|
||||
|
||||
|
||||
def compute_macro_f1(predicted, targets, num_classes=4):
|
||||
tp = torch.zeros(num_classes, device=predicted.device)
|
||||
fp = torch.zeros(num_classes, device=predicted.device)
|
||||
fn = torch.zeros(num_classes, device=predicted.device)
|
||||
for c in range(num_classes):
|
||||
tp[c] = ((predicted == c) & (targets == c)).sum()
|
||||
fp[c] = ((predicted == c) & (targets != c)).sum()
|
||||
fn[c] = ((predicted != c) & (targets == c)).sum()
|
||||
precision = tp / (tp + fp + 1e-8)
|
||||
recall = tp / (tp + fn + 1e-8)
|
||||
f1 = 2 * precision * recall / (precision + recall + 1e-8)
|
||||
return f1.mean().item()
|
||||
|
||||
|
||||
def train_one_epoch(model, train_loader, criterion, optimizer, device, epoch):
|
||||
"""训练一个epoch"""
|
||||
model.train() # 设置为训练模式
|
||||
model.train()
|
||||
running_loss = 0.0
|
||||
correct = 0
|
||||
total = 0
|
||||
all_preds, all_labels = [], []
|
||||
|
||||
# 使用 tqdm 显示进度条（可选）
|
||||
pbar = tqdm(train_loader, desc=f'Epoch {epoch + 1} [Train]')
|
||||
|
||||
for images, labels in pbar:
|
||||
# 将数据移到 GPU/CPU
|
||||
images, labels = images.to(device), labels.to(device)
|
||||
|
||||
# 前向传播
|
||||
outputs = model(images)
|
||||
loss = criterion(outputs, labels)
|
||||
|
||||
# 反向传播
|
||||
optimizer.zero_grad() # 清空梯度
|
||||
loss.backward() # 计算梯度
|
||||
optimizer.step() # 更新参数
|
||||
optimizer.zero_grad()
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
|
||||
# 统计
|
||||
running_loss += loss.item() * images.size(0)
|
||||
_, predicted = outputs.max(1)
|
||||
total += labels.size(0)
|
||||
correct += predicted.eq(labels).sum().item()
|
||||
all_preds.append(predicted)
|
||||
all_labels.append(labels)
|
||||
|
||||
# 更新进度条信息
|
||||
pbar.set_postfix({'loss': loss.item(), 'acc': 100. * correct / total})
|
||||
batch_f1 = compute_macro_f1(predicted, labels)
|
||||
pbar.set_postfix({'loss': loss.item(), 'F1': f'{batch_f1:.4f}', 'Acc': f'{100. * correct / total:.2f}%'})
|
||||
|
||||
epoch_loss = running_loss / total
|
||||
epoch_f1 = compute_macro_f1(torch.cat(all_preds), torch.cat(all_labels))
|
||||
epoch_acc = 100. * correct / total
|
||||
return epoch_loss, epoch_acc
|
||||
return epoch_loss, epoch_f1, epoch_acc
|
||||
|
||||
|
||||
def validate(model, val_loader, criterion, device):
|
||||
"""验证函数"""
|
||||
model.eval() # 设置为评估模式
|
||||
model.eval()
|
||||
running_loss = 0.0
|
||||
correct = 0
|
||||
total = 0
|
||||
all_preds, all_labels = [], []
|
||||
|
||||
with torch.no_grad(): # 不计算梯度，节省内存
|
||||
with torch.no_grad():
|
||||
for images, labels in tqdm(val_loader, desc='[Validate]'):
|
||||
images, labels = images.to(device), labels.to(device)
|
||||
|
||||
|
|
@ -69,17 +85,31 @@ def validate(model, val_loader, criterion, device):
|
|||
_, predicted = outputs.max(1)
|
||||
total += labels.size(0)
|
||||
correct += predicted.eq(labels).sum().item()
|
||||
all_preds.append(predicted)
|
||||
all_labels.append(labels)
|
||||
|
||||
epoch_loss = running_loss / total
|
||||
epoch_f1 = compute_macro_f1(torch.cat(all_preds), torch.cat(all_labels))
|
||||
epoch_acc = 100. * correct / total
|
||||
return epoch_loss, epoch_acc
|
||||
return epoch_loss, epoch_f1, epoch_acc
|
||||
|
||||
|
||||
def compute_class_weights(dataset, num_classes=4, device='cpu'):
|
||||
class_counts = torch.zeros(num_classes)
|
||||
for _, label in dataset.samples:
|
||||
lbl = label.item() if isinstance(label, torch.Tensor) else label
|
||||
class_counts[lbl] += 1
|
||||
total = class_counts.sum()
|
||||
weights = total / (num_classes * class_counts)
|
||||
return weights.to(device)
|
||||
|
||||
|
||||
def train(model, train_loader, val_loader, epochs=50, lr=0.001, device='cuda'):
|
||||
"""主训练函数"""
|
||||
|
||||
# 1. 定义损失函数和优化器
|
||||
criterion = nn.CrossEntropyLoss() # 多分类用交叉熵
|
||||
class_weights = compute_class_weights(train_loader.dataset, num_classes=4, device=device)
|
||||
criterion = nn.CrossEntropyLoss(weight=class_weights) # 多分类用交叉熵
|
||||
|
||||
# 或者使用 SGD + 动量
|
||||
optimizer = optim.SGD(model.parameters(), lr=lr, momentum=0.9, weight_decay=1e-4)
|
||||
|
|
@ -90,12 +120,12 @@ def train(model, train_loader, val_loader, epochs=50, lr=0.001, device='cuda'):
|
|||
# 2. 记录训练历史
|
||||
history = {
|
||||
'train_loss': [],
|
||||
'train_acc': [],
|
||||
'train_f1': [],
|
||||
'val_loss': [],
|
||||
'val_acc': []
|
||||
'val_f1': []
|
||||
}
|
||||
|
||||
best_val_acc = 0.0
|
||||
best_val_f1 = 0.0
|
||||
|
||||
# 3. 开始训练
|
||||
for epoch in range(epochs):
|
||||
|
|
@ -103,36 +133,36 @@ def train(model, train_loader, val_loader, epochs=50, lr=0.001, device='cuda'):
|
|||
print(f'Epoch {epoch + 1}/{epochs}')
|
||||
|
||||
# 训练
|
||||
train_loss, train_acc = train_one_epoch(model, train_loader, criterion,
|
||||
optimizer, device, epoch)
|
||||
train_loss, train_f1, train_acc = train_one_epoch(model, train_loader, criterion,
|
||||
optimizer, device, epoch)
|
||||
|
||||
# 验证
|
||||
val_loss, val_acc = validate(model, val_loader, criterion, device)
|
||||
val_loss, val_f1, val_acc = validate(model, val_loader, criterion, device)
|
||||
|
||||
# 更新学习率
|
||||
scheduler.step()
|
||||
|
||||
# 记录
|
||||
history['train_loss'].append(train_loss)
|
||||
history['train_acc'].append(train_acc)
|
||||
history['train_f1'].append(train_f1)
|
||||
history['val_loss'].append(val_loss)
|
||||
history['val_acc'].append(val_acc)
|
||||
history['val_f1'].append(val_f1)
|
||||
|
||||
# 打印结果
|
||||
print(f'Train Loss: {train_loss:.4f} | Train Acc: {train_acc:.2f}%')
|
||||
print(f'Val Loss: {val_loss:.4f} | Val Acc: {val_acc:.2f}%')
|
||||
print(f'Train Loss: {train_loss:.4f} | Train Acc: {train_acc:.2f}% | Train Macro-F1: {train_f1:.4f}')
|
||||
print(f'Val Loss: {val_loss:.4f} | Val Acc: {val_acc:.2f}% | Val Macro-F1: {val_f1:.4f}')
|
||||
print(f'Learning Rate: {optimizer.param_groups[0]["lr"]:.6f}')
|
||||
|
||||
# 保存最佳模型
|
||||
if val_acc > best_val_acc:
|
||||
best_val_acc = val_acc
|
||||
if val_f1 > best_val_f1:
|
||||
best_val_f1 = val_f1
|
||||
torch.save(model.state_dict(), 'best_model.pth')
|
||||
print(f'✓ 保存最佳模型 (Acc: {val_acc:.2f}%)')
|
||||
print(f'✓ 保存最佳模型 (Macro-F1: {val_f1:.4f})')
|
||||
|
||||
# 4. 绘制训练曲线
|
||||
|
||||
print(f'\n{"=" * 50}')
|
||||
print(f'训练完成！最佳验证准确率: {best_val_acc:.2f}%')
|
||||
print(f'训练完成！最佳验证 Macro-F1: {best_val_f1:.4f}')
|
||||
|
||||
return model, history
|
||||
|
||||
|
|
@ -153,7 +183,7 @@ if __name__ == '__main__':
|
|||
model = Net(num_classes=4) # 根据你的 Net 类调整
|
||||
#断点继续训练
|
||||
if os.path.exists('best_model.pth'):
|
||||
model.load_state_dict(torch.load('best_model.pth'))
|
||||
model.load_state_dict(torch.load('best_model.pth',map_location=torch.device('cpu')))
|
||||
model = model.to(device)
|
||||
|
||||
# 打印模型信息
|
||||
|
|
|
|||
Loading…
Reference in a new issue