# 深度学习分类框架 **Repository Path**: tantil/deep-learning-classification-framework ## Basic Information - **Project Name**: 深度学习分类框架 - **Description**: 这是一个深度学习框架，用于分类的。 - **Primary Language**: Python - **License**: Apache-2.0 - **Default Branch**: master - **Homepage**: None - **GVP Project**: No ## Statistics - **Stars**: 0 - **Forks**: 2 - **Created**: 2024-09-19 - **Last Updated**: 2024-09-20 ## Categories & Tags **Categories**: Uncategorized **Tags**: None ## README ## 数据集组织方式 假设你的数据集按照如下方式组织： ``` dataset/ ├── train/ │ ├── class_1/ │ │ ├── img1.jpg │ │ ├── img2.jpg │ │ └── ... │ ├── class_2/ │ │ ├── img1.jpg │ │ ├── img2.jpg │ │ └── ... │ └── ... ├── test/ │ ├── class_1/ │ │ ├── img1.jpg │ │ ├── img2.jpg │ │ └── ... │ ├── class_2/ │ │ ├── img1.jpg │ │ ├── img2.jpg │ │ └── ... │ └── ... └── val/ ├── class_1/ │ ├── img1.jpg │ ├── img2.jpg │ └── ... ├── class_2/ │ ├── img1.jpg │ ├── img2.jpg │ └── ... └── ... ``` ### 修改后的代码 下面是修改后的`CustomImageDataset`类： ```python import os from torch.utils.data import Dataset from torchvision import transforms from PIL import Image class CustomImageDataset(Dataset): def __init__(self, root_dir, transform=None): """ 初始化数据集类。 Args: root_dir (str): 包含所有图像的目录。 transform (callable, optional): 应用于样本的可选变换。 """ self.root_dir = root_dir self.transform = transform self.samples = self._create_samples(root_dir) def _create_samples(self, root_dir): """ 创建包含图像路径和标签的样本列表。 Returns: List[Tuple[str, int]]: 图像路径和标签组成的列表。 """ samples = [] for class_name in os.listdir(root_dir): class_dir = os.path.join(root_dir, class_name) if os.path.isdir(class_dir): for img_name in os.listdir(class_dir): img_path = os.path.join(class_dir, img_name) if os.path.isfile(img_path): samples.append((img_path, int(class_name))) return samples def __len__(self): """ 返回数据集中样本的数量。 Returns: int: 数据集中样本的数量。 """ return len(self.samples) def __getitem__(self, idx): """ 根据索引返回一个包含图像和标签的数据。 Args: idx (int): 样本的索引。 Returns: Tuple[Image, int]: 图像和对应的标签。 """ img_path, label = self.samples[idx] image = Image.open(img_path).convert('RGB') if self.transform: image = self.transform(image) return image, label ``` ### 解释 1. **`__init__` 方法**： - 初始化数据集类，设置根目录和变换。 - 调用`_create_samples`方法来创建样本列表。 2. **`_create_samples` 方法**： - 遍历根目录下的所有子目录（假设每个子目录代表一个类别）。 - 对于每个子目录，遍历其中的所有图像文件，并将其路径和类别标签添加到`samples`列表中。 3. **`__len__` 方法**： - 返回`samples`列表的长度，即数据集中样本的数量。 4. **`__getitem__` 方法**： - 根据索引`idx`从`samples`列表中获取图像路径和标签。 - 使用`PIL.Image`打开图像，并将其转换为RGB模式。 - 如果提供了变换，则应用变换。 - 返回图像和标签。 ### 使用说明 - **路径设置**：确保`root_dir`参数指向包含图像的正确目录。 - **变换设置**：根据需要设置`transform`参数，例如： ```python transform = transforms.Compose([ transforms.Resize((112, 112)), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) ``` - **类别标签**：在这个例子中，类别标签简单地由子目录名转换为整数。如果类别标签有不同的形式，需要相应调整标签处理逻辑。 。 好的，下面我将给你展示如何定义一个简单的卷积神经网络（CNN）模型，该模型满足你提到的要求：输入是一个包含batch_size信息的tensor，输出也是一个包含batch_size信息的分类结果tensor，其形状为(batch_size, num_classes)。 ## 定义一个简单的模型 我们将定义一个基础的CNN模型，该模型包含卷积层、池化层和全连接层。这个模型可以用于图像分类任务。 在定义模型的时候请注意： 模型输入仅为一个tensor，这个tensor是包含batch_size信息的 图片视频 模型输出的形状仅为一个tensor，这个tensor是包含batch_size信息的分类结果 其shape为(batch_size, num_classes) #### 示例代码 ```python import torch import torch.nn as nn import torch.nn.functional as F class SimpleCNN(nn.Module): def __init__(self, num_classes=5): super(SimpleCNN, self).__init__() # 卷积层 self.conv1 = nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, padding=1) self.conv2 = nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, padding=1) self.pool = nn.MaxPool2d(kernel_size=2, stride=2) # 全连接层 self.fc1 = nn.Linear(32 * 28 * 28, 128) # 假设输入图像大小为112x112，经过两次pooling后变为28x28 self.fc2 = nn.Linear(128, num_classes) def forward(self, x): # 输入尺寸: (batch_size, 3, 112, 112) x = F.relu(self.conv1(x)) x = self.pool(x) x = F.relu(self.conv2(x)) x = self.pool(x) # 在送入全连接层之前，需要展平张量 x = x.view(-1, 32 * 28 * 28) x = F.relu(self.fc1(x)) x = self.fc2(x) return x # 创建模型实例 model = SimpleCNN(num_classes=5) print(model) ``` ### 说明 1. **模型定义**： - `SimpleCNN`继承自`nn.Module`。 - 构造函数`__init__`中定义了卷积层、池化层和全连接层。 - `forward`方法定义了数据流经网络的过程。 2. **卷积层**： - 第一个卷积层接收3通道的输入图像（通常是RGB图像），输出16个特征图。 - 第二个卷积层接收上一层的输出作为输入，输出32个特征图。 3. **池化层**： - 使用最大池化层减少空间维度。 4. **全连接层**： - 在送入全连接层之前，需要将特征图展平成一维向量。 - 第一个全连接层之后使用ReLU激活函数。 - 输出层不使用激活函数，因为通常会在损失函数中应用softmax。 ## 如何使用模型 ### main.py的修改 ### 1. 修改数据集路径 - **路径更改**：如果数据集存储的位置不同，需要更改`root_dir`参数。例如，如果你的训练数据存储在`/path/to/train`目录下，那么应该将`root_dir='dataset/train'`改为`root_dir='/path/to/train'`。 ### 2. 自定义数据集类 - **数据集类**：确保`CustomImageDataset`类正确实现了`__init__`, `__len__`, 和 `__getitem__`方法。如果你需要处理的数据格式不同（如视频帧），可能需要重写这个类。 - **文件路径**：如果数据集的组织方式不同，需要修改`CustomImageDataset`类中读取图像的方式。 ### 3. 修改图像转换 - **变换操作**：如果你需要不同的图像预处理步骤，可以在`transforms.Compose`中添加或删除变换操作。例如，可以增加`transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])`来进行标准化。 ### 4. 自定义模型 - **模型定义**：如果你有自己的模型定义，需要确保这些模型类已经在代码中正确导入。例如，如果`OurModel`是你的自定义模型，确保`OthersModels.py`文件中已经定义了这个类。 - **模型参数**：在`models`列表中定义模型时，可以根据需要调整参数，例如`num_classes`应与你的分类任务匹配。 ### 5. 修改训练和测试过程 - **训练函数**：检查`my_train`和`my_test`函数，了解它们是如何工作的。如果需要，可以修改这些函数以适应不同的训练策略或评估指标。 - **批量大小**：可以根据GPU内存调整`batch_size`。更大的批量可能会导致内存不足错误，而较小的批量可能会延长训练时间。 ### 6. 添加新的模型 - **模型添加**：如果想要添加一个新的模型到训练过程中，只需在`models`列表中加入新的模型实例化字符串即可。例如： ```python models.append('SimpleCNN(num_classes=5)') ``` ### 7. 处理预训练权重 - **权重检查**：如果希望在训练前加载预训练权重，可以在训练之前检查是否存在预训练的.pth文件，并加载这些权重。 - **权重保存**：训练完成后，通常会希望保存模型权重。可以修改代码，在训练结束后保存模型状态字典。 ### 示例修改 假设你想添加一个新的自定义模型`MyCustomModel`，并且希望在训练前加载预训练权重，并在训练后保存模型权重，可以这样修改： ```python from OthersModels import SimpleCNN # 假设这是你的自定义模型 models.append('SimpleCNN(num_classes=5)') for model_str in models: model = eval(model_str) print("") print("开始训练模型：", model.model_name) # 如果 模型参数文件夹下有 model.model_name.pth文件则加载预训练权重 if os.path.exists(f"模型参数/{model.model_name}.pth"): model.load_state_dict(torch.load(f"模型参数/{model.model_name}.pth")) print("模型参数已加载") model.my_train(epochs=20, data_loader=data_loader_train, data_loader_eval=data_loader_eval) model.my_test(data_loader_eval) # 保存模型权重 torch.save(model.state_dict(), f"模型参数/{model.model_name}.pth") # 清除显存 torch.cuda.empty_cache() ``` ## 剩下的内容请自己阅读代码