【机器学习】imagenet2012 数据预处理数据预处理

伪NChris2025-03-27 11:53

【机器学习】数据预处理

  • [1. 下载/解压数据](#1. 下载/解压数据)
  • [2. 数据预处理](#2. 数据预处理)
  • [3. 加载以及训练代码](#3. 加载以及训练代码)
    • [3.1 使用PIL等加载代码](#3.1 使用PIL等加载代码)
    • [3.2 使用OpenCV的方式来一张张加载代码](#3.2 使用OpenCV的方式来一张张加载代码)
    • [3.3 h5的方式来加载大文件](#3.3 h5的方式来加载大文件)
  • 最后总结

这个数据大约 140个G,128w的训练集

1. 下载/解压数据

首先需要下载数据：

数据最后处理成如图的格式，每个种类的数据都放到一个相同的文件夹中去，这里的文件夹名称（种类名称）最好改成整数，方便后续处理

2. 数据预处理

需要对数据做如下处理

  1. 处理成模型需要的224*224长宽的数据
  2. 处理成h5/npy之类大文件格式，从而减少CPU的IO开支
python 复制代码 
import h5py
import numpy as np
import os
from tqdm import tqdm
import cv2
from concurrent.futures import ThreadPoolExecutor
from sklearn.preprocessing import LabelEncoder

def process_image(file_path, size=(224, 224)):
    image = cv2.imread(file_path)
    if image is None:
        print(f"无法读取图像: {file_path}")
        return None

    # 调整图像大小
    resized_image = cv2.resize(image, size)
    return resized_image

def create_hdf5_datasets(input_dir, output_dir, images_per_file=1000, max_workers=8):
    # 获取所有文件的列表
    all_files = []
    for root, dirs, files in os.walk(input_dir):
        for file_name in files:
            file_path = os.path.join(root, file_name)
            all_files.append(file_path)

    # 确保输出目录存在
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    # 获取所有标签并进行编码
    all_labels = [os.path.basename(os.path.dirname(file)) for file in all_files]
    label_encoder = LabelEncoder()
    label_encoder.fit(all_labels)

    file_count = 0
    total_files = len(all_files)

    # 使用多线程处理图像
    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        for i in range(0, total_files, images_per_file):
            chunk_files = all_files[i:i + images_per_file]
            processed_images = list(tqdm(executor.map(process_image, chunk_files), total=len(chunk_files), desc=f"Processing chunk {file_count + 1}"))

            # 过滤掉 None 值
            processed_images = [img for img in processed_images if img is not None]

            # 创建标签数据（假设标签为文件夹名称）
            labels = [os.path.basename(os.path.dirname(file)) for file in chunk_files if cv2.imread(file) is not None]
            encoded_labels = label_encoder.transform(labels)

            # 写入 HDF5 文件
            output_hdf5 = os.path.join(output_dir, f'train_{file_count + 1}.hdf5')
            with h5py.File(output_hdf5, 'w') as f:
                dataset_images = f.create_dataset("images", (len(processed_images), 224, 224, 3), dtype='uint8')
                dataset_labels = f.create_dataset("labels", (len(encoded_labels),), dtype='int')
                for j, img in enumerate(processed_images):
                    dataset_images[j] = img
                    dataset_labels[j] = encoded_labels[j]

            file_count += 1
            print(f"Created {output_hdf5} with {len(processed_images)} images")

    print(f"Total HDF5 files created: {file_count}")

# 示例用法
input_directory_path = 'E:\\data\\train'  # 替换为你的目录路径  
output_directory_path = 'E:\\data\\hdf5\\train'  # 输出的目录路径
create_hdf5_datasets(input_directory_path, output_directory_path, images_per_file=50000)  # 创建多个 HDF5 文件

这里就是将图片分成若干份，每一份50000张图，主要是我电脑内存32G 无法一次性加载，所以分割了一下。

3. 加载以及训练代码

3.1 使用PIL等加载代码

这个方式是一张张的加载图片，加载后再处理成模型需要的尺寸，在一张张加载图片的时候速度较慢，会影响训练速度

python 复制代码 
# 定义自定义数据集类
class CustomDataset(torch.utils.data.Dataset):
    def __init__(self, csv_file, transform=None):
        self.data_frame = pd.read_csv(csv_file)
        self.transform = transform
        self.label_encoder = LabelEncoder()
        self.data_frame['label'] = self.label_encoder.fit_transform(self.data_frame['label'])  # 将标签编码为整数

    def __len__(self):
        return len(self.data_frame)

    def __getitem__(self, idx):
        img_path = self.data_frame.iloc[idx, 0] # 图像路径
        image = Image.open(train_file + img_path).convert('RGB')# 读取图像
        label = self.data_frame.iloc[idx, 1] #从表格中读取标签 ，此时标签已经被编码为整数
        label = torch.tensor(label, dtype=torch.long)# 将标签转换为张量

        if self.transform:
            image = self.transform(image)

        return image, label

# 定义图像转换
transform = transforms.Compose([
    transforms.Resize((224, 224)),  # 调整图像大小
    transforms.ToTensor(),          # 转换为张量
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # 归一化 mean的值和std的值是根据ImageNet数据集的均值和标准差计算得到的
])

3.2 使用OpenCV的方式来一张张加载代码

OpenCV确实能加速一点IO的速度，

python 复制代码 
import os
import pandas as pd
import cv2  # 导入 OpenCV 库
from sklearn.preprocessing import LabelEncoder
import torch
from torch.utils.data import Dataset, DataLoader, random_split
import torchvision.transforms as transforms
import torchvision.models as models
import torch.nn as nn
import torch.optim as optim
from tqdm import tqdm  # 导入 tqdm 库
import time

# 定义数据路径
data_path = 'E:\\data\\ImageNet2012\\ILSVRC2012_img_train\\'

# 定义自定义数据集类
class CustomDataset(torch.utils.data.Dataset):
    def __init__(self, csv_file, data_path, transform=None):
        self.data_frame = pd.read_csv(csv_file)
        self.data_path = data_path
        self.transform = transform
        self.label_encoder = LabelEncoder()
        self.data_frame['label'] = self.label_encoder.fit_transform(self.data_frame['label'])  # 将标签编码为整数

    def __len__(self):
        return len(self.data_frame)

    def __getitem__(self, idx):
        start_time = time.time()
        data_load_time = time.time() - start_time

        img_name = self.data_frame.iloc[idx, 0]  # 图像相对路径
        img_path = os.path.join(self.data_path, img_name)  # 生成完整的图像路径
        image = cv2.imread(img_path)  # 使用 OpenCV 读取图像
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)  # 将图像从 BGR 转换为 RGB
        image = cv2.resize(image, (224, 224))  # 调整图像大小
        label = self.data_frame.iloc[idx, 1]  # 从表格中读取标签，此时标签已经被编码为整数
        label = torch.tensor(label, dtype=torch.long)  # 将标签转换为张量

        data_to_device_time = time.time() - start_time - data_load_time

        if self.transform:
            image = self.transform(image)

        forward_time = time.time() - start_time - data_load_time - data_to_device_time
        print(f"Data load time: {data_load_time:.4f}, Data to device time: {data_to_device_time:.4f}, Forward time: {forward_time:.4f}")
        return image, label

# 定义图像转换
transform = transforms.Compose([
    transforms.ToTensor(),          # 转换为张量
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # 归一化 mean的值和std的值是根据ImageNet数据集的均值和标准差计算得到的
])

# 创建数据集
csv_file = os.path.join(data_path, 'train.csv')
dataset = CustomDataset(csv_file=csv_file, data_path=data_path, transform=transform)

# 将数据集分为训练集和验证集
train_size = int(0.8 * len(dataset))
val_size = len(dataset) - train_size
train_dataset, val_dataset = random_split(dataset, [train_size, val_size])

# 创建数据加载器
train_dataloader = DataLoader(train_dataset, batch_size=512, shuffle=True)  # 设置 shuffle 为 True
val_dataloader = DataLoader(val_dataset, batch_size=512, shuffle=False) 

# 加载预训练的 ResNet 模型
model = models.resnet18(pretrained=True)
num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs, len(dataset.data_frame['label'].unique()))  # 根据标签数量调整最后一层

# 将模型移动到 GPU 上
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

# 训练函数
def train_model(model, dataloader, criterion, optimizer):
    model.train()
    running_loss = 0.0
    for inputs, labels in tqdm(dataloader, desc="Training"):  # 使用 tqdm 包装 dataloader
        inputs, labels = inputs.to(device), labels.to(device)  # 将数据移动到 GPU 上
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item() * inputs.size(0)
    epoch_loss = running_loss / len(dataloader.dataset)
    print(f'Training Loss: {epoch_loss:.4f}')


# 测试函数
def test_model(model, dataloader, criterion):
    model.eval()
    correct = 0
    total = 0
    running_loss = 0.0
    with torch.no_grad():
        for inputs, labels in tqdm(dataloader, desc="Validation"):  # 使用 tqdm 包装 dataloader
            inputs, labels = inputs.to(device), labels.to(device)  # 将数据移动到 GPU 上
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            running_loss += loss.item() * inputs.size(0)
            _, predicted = torch.max(outputs, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    accuracy = correct / total
    epoch_loss = running_loss / len(dataloader.dataset)
    print(f'Test Loss: {epoch_loss:.4f}, Accuracy: {accuracy:.4f}')

# 训练和验证循环
epochs = 25
for t in range(epochs):
    print(f"Epoch {t+1}\n-------------------------------")
    train_model(model, train_dataloader, criterion, optimizer)
    print("Validation:")
    test_model(model, val_dataloader, criterion)
print("Done!")

3.3 h5的方式来加载大文件

HDF5Dataset 类在初始化时只加载文件索引，而不是加载所有数据。在 getitem 方法中，它会根据索引动态加载所需的 HDF5 文件，并从中读取图像和标签。这可以确保在每次访问数据时只加载当前需要的 HDF5 文件，并在使用完成后自动从内存中移除。

python 复制代码 
import os
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset
from torchvision import models, transforms
from tqdm import tqdm
import h5py

# 定义数据路径
train_data_path = 'E:\\data\\hdf5\\train'
val_data_path = 'E:\\data\\hdf5\\val'

# 定义自定义数据集类
class HDF5Dataset(Dataset):
    def __init__(self, hdf5_dir, transform=None):
        self.hdf5_files = [os.path.join(hdf5_dir, f) for f in os.listdir(hdf5_dir) if f.endswith('.hdf5')]
        self.transform = transform
        self.file_indices = []
        self.load_file_indices()

    def load_file_indices(self):
        for file_idx, hdf5_file in enumerate(self.hdf5_files):
            with h5py.File(hdf5_file, 'r') as f:
                num_images = f['images'].shape[0]
                self.file_indices.extend([(file_idx, i) for i in range(num_images)])

    def __len__(self):
        return len(self.file_indices)

    def __getitem__(self, idx):
        file_idx, image_idx = self.file_indices[idx]
        hdf5_file = self.hdf5_files[file_idx]

        with h5py.File(hdf5_file, 'r') as f:
            image = f['images'][image_idx]
            label = f['labels'][image_idx]

        if self.transform:
            image = self.transform(image)

        # 将标签转换为张量
        label = torch.tensor(label, dtype=torch.long)

        return image, label

# 定义图像转换
transform = transforms.Compose([
    transforms.ToTensor(),          # 转换为张量
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # 归一化 mean的值和std的值是根据ImageNet数据集的均值和标准差计算得到的
])

# 创建训练集数据集
train_dataset = HDF5Dataset(hdf5_dir=train_data_path, transform=transform)

# 创建验证集数据集
val_dataset = HDF5Dataset(hdf5_dir=val_data_path, transform=transform)

# 创建数据加载器
train_dataloader = DataLoader(train_dataset, batch_size=256, shuffle=True)  # 设置 shuffle 为 True
val_dataloader = DataLoader(val_dataset, batch_size=256, shuffle=False) 

# 加载预训练的 ResNet 模型
model = models.resnet18(pretrained=True)
num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs, len(set(train_dataset.file_indices)))  # 根据标签数量调整最后一层

# 将模型移动到 GPU 上
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

# 训练函数
def train_model(model, dataloader, criterion, optimizer):
    model.train()
    running_loss = 0.0
    for inputs, labels in tqdm(dataloader, desc="Training"):  # 使用 tqdm 包装 dataloader
        inputs, labels = inputs.to(device), labels.to(device)  # 将数据移动到 GPU 上
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item() * inputs.size(0)
    epoch_loss = running_loss / len(dataloader.dataset)
    print(f'Training Loss: {epoch_loss:.4f}')

# 测试函数
def test_model(model, dataloader, criterion):
    model.eval()
    correct = 0
    total = 0
    running_loss = 0.0
    with torch.no_grad():
        for inputs, labels in tqdm(dataloader, desc="Validation"):  # 使用 tqdm 包装 dataloader
            inputs, labels = inputs.to(device), labels.to(device)  # 将数据移动到 GPU 上
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            running_loss += loss.item() * inputs.size(0)
            _, predicted = torch.max(outputs, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    accuracy = correct / total
    epoch_loss = running_loss / len(dataloader.dataset)
    print(f'Test Loss: {epoch_loss:.4f}, Accuracy: {accuracy:.4f}')

# 训练和验证循环
epochs = 25
model_save_path = 'model_checkpoint.pth'
for t in range(epochs):
    print(f"Epoch {t+1}\n-------------------------------")
    train_model(model, train_dataloader, criterion, optimizer)
    print("Validation:")
    test_model(model, val_dataloader, criterion)
    
    # 每5个循环保存一次模型，并删除之前的模型
    if (t + 1) % 5 == 0:
        if os.path.exists(model_save_path):
            os.remove(model_save_path)
        torch.save(model.state_dict(), model_save_path)
        print(f"Model saved at epoch {t+1}")

print("Done!")

最后总结

我的电脑环境 i5 12400+4090+32G内存+固态。

但磁盘速度才几十M如果是机械盘的话应该也问题不大

然后我训练的时间最快能达到45分钟一个epoch，使用3.3章节中的代码。

提升训练速度的小技巧

  1. 不要开任务管理器，虽然开着很爽，但确实比较占用CPU的资源
  2. 不要开浏览器，浏览器中不知道运行了些什么东西会影响速度
  3. 不要开很多vscode，只保留debug的一个，能加速10分钟
相关推荐
鬼手爱剪辑2 小时前
短剧AI 配音技术的难点、原理、开源和闭源软件汇总-2026 年
人工智能
田里的水稻2 小时前
PyTorch_Dataloader\TensorBoard---数据加载与可视化
人工智能·pytorch·python·神经网络·机器学习
卿卿的产品经理日记2 小时前
【AI产品经理实战】Day 38｜try-except 踩坑实录：报错是数据
人工智能·笔记·python·学习·产品经理
cxj75142 小时前
GEO 工程化实践：关键页面的快慢双轨更新与多引擎引用维护
人工智能·geo·geo优化
Java_AI工程师2 小时前
Spring4 / SpringBoot1.x 老项目想接入 AI？不用升级框架！3 套无痛落地方案
java·人工智能·spring boot
znnnk2 小时前
【AI技术】Vibe Coding 实践指南：从项目规划到高效开发
人工智能·ai·ai编程·ai ide·vibe coding
Simon_lca2 小时前
ZARA、H&M、SHEIN验厂差异全对比：一套底层记录，三套接口输出
大数据·人工智能·经验分享·科技·制造
爱喝雪碧的可乐2 小时前
DeepSeek 崩了近5小时：一次故障扒透「单点依赖、重试风暴与大模型高可用」
大数据·人工智能·deepseek
GlobalTradercrb2 小时前
独立站SEO先写博客还是先优化产品页？看站点阶段
数据库·人工智能·google·pdf·go·seo
田里的水稻3 小时前
FA_CNN卷积神经网络三---非线性激活函数PyTorch
人工智能·pytorch·神经网络·机器学习·cnn
热门推荐
01国内可直接用、免费额度/永久免费的大模型API清单（含 SiliconFlow、火山、阿里、智谱、百度、Kimi、DeepSeek、DMXAPI 等）02如何新建文件夹？ 电脑新建文件夹的4种方法03历年考研数学一、数学二、数学三真题试卷及答案PDF04GitHub 镜像站点05「全语录+省流版」邵艾伦 × 孙宇晨：聊风口、聊个人IP、聊年轻人如何抓住 AI 时代的机会06几个好用的ip纯净度检测网站07苹果iPA游戏软件资源下载网站08昇腾910B系列芯片完全指南：四款型号对比、服务器选型与大模型部署实战09Windows 10/11 如何退出微软账户并改用本地账户登录（含找不到选项的解决方法）10乙二醇水溶液物性参数