如何训练自己的数据集之——恶劣天气目标检测数据集

 weather_detection_dataset/

    
 ├── images/
    
 │   ├── train/
    
 │   │   ├── train_image_0001.jpg
    
 │   │   ├── train_image_0002.jpg
    
 │   │   └── ...
    
 │   ├── val/
    
 │   │   ├── val_image_0001.jpg
    
 │   │   ├── val_image_0002.jpg
    
 │   │   └── ...
    
 ├── labels_yolo/
    
 │   ├── train/
    
 │   │   ├── train_image_0001.txt
    
 │   │   ├── train_image_0002.txt
    
 │   │   └── ...
    
 │   ├── val/
    
 │   │   ├── val_image_0001.txt
    
 │   │   ├── val_image_0002.txt
    
 │   │   └── ...
    
 ├── labels_voc/
    
 │   ├── train/
    
 │   │   ├── train_image_0001.xml
    
 │   │   ├── train_image_0002.xml
    
 │   │   └── ...
    
 │   ├── val/
    
 │   │   ├── val_image_0001.xml
    
 │   │   ├── val_image_0002.xml
    
 │   │   └── ...
    
 └── data.yaml  # 数据集配置文件
    
    
    
    
    
![](https://ad.itadn.com/c/weblog/blog-img/images/2025-07-13/VvEiHhJx8yYeuK9Zf5mFNqWdUrgL.png)

 # data.yaml 文件

    
 train: ../weather_detection_dataset/images/train/
    
 val: ../weather_detection_dataset/images/val/
    
 test: ../weather_detection_dataset/images/test/  # optional
    
  
    
 nc: 8  # number of classes
    
 names: ['person', 'rider', 'car', 'truck', 'bus', 'train', 'motorcycle', 'bicycle']  # class names

 git clone https://github.com/ultralytics/yolov5.git  # clone repo

    
 cd yolov5
    
 pip install -r requirements.txt  # install dependencies

 cd yolov5

    
 python train.py --img 640 --batch 16 --epochs 300 --data ../weather_detection_dataset/data.yaml --weights yolov5s.pt --cache

 import torch

    
 from utils.datasets import LoadImagesAndLabels  # 导入数据加载器
    
 from models.experimental import attempt_load  # 导入模型加载器
    
 from utils.torch_utils import select_device  # 导入选用设备的函数
    
 from utils.general import check_dataset  # 导入检查数据集的函数
    
  
    
 def main():
    
     device = select_device('')  # 选择设备，自动选择GPU/CPU
    
     data_yaml = '../weather_detection_dataset/data.yaml'
    
     train_images_folder = '../weather_detection_dataset/images/train/'
    
     train_labels_folder = '../weather_detection_dataset/labels_yolo/train/'
    
     val_images_folder = '../weather_detection_dataset/images/val/'
    
     val_labels_folder = '../weather_detection_dataset/labels_yolo/val/'
    
  
    
     # 加载数据集
    
     train_set = LoadImagesAndLabels(train_images_folder, train_labels_folder)
    
     val_set = LoadImagesAndLabels(val_images_folder, val_labels_folder)
    
  
    
     # 创建数据加载器
    
     train_loader = torch.utils.data.DataLoader(train_set, batch_size=16, shuffle=True, num_workers=4)
    
     val_loader = torch.utils.data.DataLoader(val_set, batch_size=16, shuffle=False, num_workers=4)
    
  
    
     # 加载预训练模型
    
     model = attempt_load('yolov5s.pt', map_location=device)  # 加载预训练权重
    
     model.to(device)
    
  
    
     # 设置损失函数和优化器
    
     criterion = torch.nn.CrossEntropyLoss()
    
     optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
    
  
    
     # 训练循环
    
     epochs = 300
    
     for epoch in range(epochs):
    
     model.train()
    
     for images, targets, _, _ in train_loader:
    
         images = images.to(device)
    
         targets = [t.to(device) for t in targets]
    
  
    
         # 前向传播
    
         outputs = model(images)
    
  
    
         # 计算损失
    
         loss = criterion(outputs, targets)
    
  
    
         # 反向传播和优化
    
         optimizer.zero_grad()
    
         loss.backward()
    
         optimizer.step()
    
  
    
     print(f'Epoch [{epoch+1}/{epochs}], Loss: {loss.item()}')
    
  
    
     torch.save(model.state_dict(), 'trained_model.pt')
    
     print('Training complete.')
    
  
    
 if __name__ == '__main__':
    
     main()
    
    
    
    
    
![](https://ad.itadn.com/c/weblog/blog-img/images/2025-07-13/Vdg4rEaItNPljD6K8cHqbZRGkXBz.png)

    python detect.py --weights runs/train/exp/weights/best.pt --img 640 --conf 0.4 --source ../weather_detection_dataset/images/test/