yukun-hh/trash-division
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

重置所有文件与main一致，新增app.py（Gradio推理前端）

Browse Source
This commit is contained in:
weikaiwen348-code
2026-06-03 20:33:22 +08:00
parent 543e833fd0
commit 0e9bd58dce
24 changed files with 1957 additions and 168 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Train.py
+87 -68
View File
@@ -12,52 +12,70 @@ import torch.optim as optim
from tqdm import tqdm # 进度条，可选
import matplotlib.pyplot as plt
from Model import Net
from Dataloader import create_dataloaders
import os
import csv
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)
@@ -68,37 +86,52 @@ 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) # 多分类用交叉熵
# 优化器选择（推荐 Adam 或 SGD）
optimizer = optim.Adam(model.parameters(), lr=lr, weight_decay=1e-4)
# 或者使用 SGD + 动量
# optimizer = optim.SGD(model.parameters(), lr=lr, momentum=0.9, weight_decay=1e-4)
optimizer = optim.SGD(model.parameters(), lr=lr, momentum=0.9, weight_decay=1e-4)
# 学习率调度器（可选，帮助收敛）
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=20, gamma=0.1)
# 或者用余弦退火
# scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)
scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)
# 2. 记录训练历史
history = {
'train_loss': [],
'train_f1': [],
'train_acc': [],
'val_loss': [],
'val_f1': [],
'val_acc': []
}
best_val_acc = 0.0
best_val_f1 = 0.0
log_file = open('training_log.csv', 'w', newline='')
log_writer = csv.writer(log_file)
log_writer.writerow(['epoch', 'train_loss', 'train_f1', 'train_acc', 'val_loss', 'val_f1', 'val_acc', 'lr', 'best'])
# 3. 开始训练
for epoch in range(epochs):
@@ -106,77 +139,65 @@ 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_f1'].append(train_f1)
history['train_acc'].append(train_acc)
history['val_loss'].append(val_loss)
history['val_f1'].append(val_f1)
history['val_acc'].append(val_acc)
# 打印结果
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
best_mark = ''
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}%)')
best_mark = 'best'
print(f'✓ 保存最佳模型 (Macro-F1: {val_f1:.4f})')
lr = optimizer.param_groups[0]['lr']
log_writer.writerow([epoch + 1, train_loss, train_f1, train_acc, val_loss, val_f1, val_acc, lr, best_mark])
log_file.flush()
# 4. 绘制训练曲线
plot_training_history(history)
print(f'\n{"=" * 50}')
print(f'训练完成！最佳验证准确率: {best_val_acc:.2f}%')
print(f'训练完成！最佳验证 Macro-F1: {best_val_f1:.4f}')
return model, history
def plot_training_history(history):
"""绘制训练曲线"""
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))
# 损失曲线
ax1.plot(history['train_loss'], label='Train Loss')
ax1.plot(history['val_loss'], label='Val Loss')
ax1.set_xlabel('Epoch')
ax1.set_ylabel('Loss')
ax1.set_title('Training and Validation Loss')
ax1.legend()
ax1.grid(True)
# 准确率曲线
ax2.plot(history['train_acc'], label='Train Acc')
ax2.plot(history['val_acc'], label='Val Acc')
ax2.set_xlabel('Epoch')
ax2.set_ylabel('Accuracy (%)')
ax2.set_title('Training and Validation Accuracy')
ax2.legend()
ax2.grid(True)
plt.tight_layout()
plt.savefig('training_history.png', dpi=150)
plt.show()
# ========== 使用示例 ==========
if __name__ == '__main__':
# 假设你的 dataloader 已经写好了
# train_loader = ...
# val_loader = ...
train_loader, val_loader, class_names = create_dataloaders(
data_root='../trash_division_data/ultimate_4_class/', # 与trash-division同级文件夹
batch_size=16, # 根据你的显存调整
image_size=256, # 与你模型输入一致
num_workers=8, # Windows 可能需设为 0
augment=True # 训练时使用数据增强
)
# 1. 创建模型
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = Net().get_network() # 根据你的 Net 类调整
device = torch.device('cuda' if torch.cuda.is_available() else 'xpu' if torch.xpu.is_available() else 'cpu')
model = Net(num_classes=4) # 根据你的 Net 类调整
#断点继续训练
if os.path.exists('best_model.pth'):
model.load_state_dict(torch.load('best_model.pth',map_location=torch.device('cpu')))
model = model.to(device)
# 打印模型信息
@@ -188,11 +209,9 @@ if __name__ == '__main__':
model=model,
train_loader=train_loader,
val_loader=val_loader,
epochs=50,
epochs=20,
lr=0.001,
device=device
)
# 3. 加载最佳模型用于预测
model.load_state_dict(torch.load('best_model.pth'))
print('训练完成，最佳模型已加载')
model.load_state_dict(torch.load('best_model.pth'))