From b2f7a9c172ee644c951ab473c695c6fdb8287839 Mon Sep 17 00:00:00 2001 From: yukun-hh Date: Thu, 16 Apr 2026 13:55:02 +0800 Subject: [PATCH] =?UTF-8?q?=E6=A8=A1=E5=9E=8B=E8=A7=84=E8=8C=83=E8=AE=BE?= =?UTF-8?q?=E8=AE=A1?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- Dataloader.py | 4 +- Model.py | 109 +++++++++++++++++++++++++++----------------------- Train.py | 37 ++--------------- 3 files changed, 64 insertions(+), 86 deletions(-) diff --git a/Dataloader.py b/Dataloader.py index efd73c9..19b060a 100644 --- a/Dataloader.py +++ b/Dataloader.py @@ -158,9 +158,9 @@ def visualize_batch(dataloader, class_names, num_images=8): if __name__ == '__main__': train_loader, val_loader, class_names = create_dataloaders( data_root='../trash_division_data/ultimate_4_class/', # 与trash-division同级文件夹 - batch_size=32, # 根据你的显存调整 + batch_size=16, # 根据你的显存调整 image_size=256, # 与你模型输入一致 - num_workers=4, # Windows 可能需设为 0 + num_workers=16, # Windows 可能需设为 0 augment=True # 训练时使用数据增强 ) visualize_batch(train_loader, class_names, num_images=8) diff --git a/Model.py b/Model.py index d568b97..313b2d0 100644 --- a/Model.py +++ b/Model.py @@ -6,31 +6,32 @@ author : yukun-hh date : 2026-4-10 """ -#神经网络模型库 import torch from torch import nn from torch.nn import functional as F from torchsummary import summary -#残差块 -class Resblock(nn.Module): - def __init__(self, input_channels,output_channels,use_1x1conv=False,strides=1): - """ + +# 残差块 +class Resblock(nn.Module): + def __init__(self, input_channels, output_channels, use_1x1conv=False, strides=1): + """ :param input_channels: 进入残差块时的原通道 :param output_channels: 输出时的通道数 :param use_1x1conv: 如果输入和输出通道不相等时，需要用一个1x1的卷积层对原来的输入进行一个通道提升 :param strides: 默认1，如果大于1起到缩小张量的作用 """ super().__init__() - self.conv1 = nn.Conv2d(input_channels,output_channels,kernel_size=3,padding=1,stride=strides) - self.conv2 = nn.Conv2d(output_channels,output_channels,kernel_size=3,padding=1,stride=1) + self.conv1 = nn.Conv2d(input_channels, output_channels, kernel_size=3, padding=1, stride=strides) + self.conv2 = nn.Conv2d(output_channels, output_channels, kernel_size=3, padding=1, stride=1) if use_1x1conv: - self.conv3 = nn.Conv2d(input_channels, output_channels,kernel_size=1, stride=strides) + self.conv3 = nn.Conv2d(input_channels, output_channels, kernel_size=1, stride=strides) else: self.conv3 = None self.bn1 = nn.BatchNorm2d(output_channels) self.bn2 = nn.BatchNorm2d(output_channels) - def forward(self,X): + + def forward(self, X): Y = F.relu(self.bn1(self.conv1(X))) Y = self.bn2(self.conv2(Y)) if self.conv3 is not None: @@ -38,36 +39,41 @@ class Resblock(nn.Module): Y += X return F.relu(Y) -class Net(): + +class Net(nn.Module): """ 模型的主要结构就在这里了，到时也好该和调用 现在必须实现的方法： 目前还是以图片缩放到256＊256构建残差块 """ - net = nn.Sequential() - def resnet_block(self,input_channels, num_channels, num_residuals, - first_block=False): - """ - :param input_channels: 输入维度 - :param num_channels: 输出维度 - :param num_residuals: 单个残差层的残差块数 - :param first_block: 第一块不用下采样 特殊控制 - :return: list[nn.Module] - """ - blk = [] - for i in range(num_residuals): - if i == 0 and not first_block: - blk.append(Resblock(input_channels, num_channels, - use_1x1conv=True, strides=2)) - else: - blk.append(Resblock(num_channels, num_channels)) - return blk def __init__(self): - b1 = nn.Sequential( nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3), - nn.BatchNorm2d(64), nn.ReLU(), - nn.MaxPool2d(kernel_size=3, stride=2, padding=1) - ) + super().__init__() + + # 定义残差块的辅助方法 + def resnet_block(input_channels, num_channels, num_residuals, first_block=False): + """ + :param input_channels: 输入维度 + :param num_channels: 输出维度 + :param num_residuals: 单个残差层的残差块数 + :param first_block: 第一块不用下采样 特殊控制 + :return: list[nn.Module] + """ + blk = [] + for i in range(num_residuals): + if i == 0 and not first_block: + blk.append(Resblock(input_channels, num_channels, use_1x1conv=True, strides=2)) + else: + blk.append(Resblock(num_channels, num_channels)) + return blk + + # 构建网络各层 + self.b1 = nn.Sequential( + nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3), + nn.BatchNorm2d(64), + nn.ReLU(), + nn.MaxPool2d(kernel_size=3, stride=2, padding=1) + ) """ 7×7 卷积层，输出通道 64，步长 2，填充 3 (3×256×256)->(64×128×128) @@ -75,25 +81,28 @@ class Net(): 最大池化 (64×128×128)->(64×64×64) """ - b2 = nn.Sequential(*self.resnet_block(64, 64, num_residuals=3, first_block=True)) - b3 = nn.Sequential(*self.resnet_block(64, 128, num_residuals=4)) - b4 = nn.Sequential(*self.resnet_block(128, 256, num_residuals=6)) - b5 = nn.Sequential(*self.resnet_block(256, 512, num_residuals=3)) - self.net = nn.Sequential(b1, b2, b3, b4, b5,nn.AdaptiveAvgPool2d((1,1)),nn.Flatten(), nn.Linear(512, 4)) - def get_network(self): - return self.net + self.b2 = nn.Sequential(*resnet_block(64, 64, num_residuals=3, first_block=True)) + self.b3 = nn.Sequential(*resnet_block(64, 128, num_residuals=4)) + self.b4 = nn.Sequential(*resnet_block(128, 256, num_residuals=6)) + self.b5 = nn.Sequential(*resnet_block(256, 512, num_residuals=3)) + self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) + self.flatten = nn.Flatten() + self.fc = nn.Linear(512, 4) + + def forward(self, x): + x = self.b1(x) + x = self.b2(x) + x = self.b3(x) + x = self.b4(x) + x = self.b5(x) + x = self.avgpool(x) + x = self.flatten(x) + x = self.fc(x) + return x if __name__ == '__main__': - Net_new = Net() - X = torch.rand(size=(1, 3, 256, 256)) - summary(Net_new.get_network(), input_size=(3, 256, 256)) - - - - - - - - + model = Net() + # 使用 torchsummary 查看模型结构 + summary(model, input_size=(3, 256, 256)) \ No newline at end of file diff --git a/Train.py b/Train.py index e5d462c..6a71ae2 100644 --- a/Train.py +++ b/Train.py @@ -81,14 +81,11 @@ def train(model, train_loader, val_loader, epochs=50, lr=0.001, device='cuda'): criterion = nn.CrossEntropyLoss() # 多分类用交叉熵 # 优化器选择（推荐 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 = { @@ -133,7 +130,6 @@ def train(model, train_loader, val_loader, epochs=50, lr=0.001, device='cuda'): print(f'✓ 保存最佳模型 (Acc: {val_acc:.2f}%)') # 4. 绘制训练曲线 - plot_training_history(history) print(f'\n{"=" * 50}') print(f'训练完成！最佳验证准确率: {best_val_acc:.2f}%') @@ -141,33 +137,6 @@ def train(model, train_loader, val_loader, epochs=50, lr=0.001, device='cuda'): 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 已经写好了 @@ -181,7 +150,7 @@ if __name__ == '__main__': # 1. 创建模型 device = torch.device('cuda' if torch.cuda.is_available() else 'xpu' if torch.xpu.is_available() else 'cpu') - model = Net().get_network() # 根据你的 Net 类调整 + model = Net() # 根据你的 Net 类调整 model = model.to(device) # 打印模型信息