基于.mat文件构建PyTorch自定义Dataset类的实现方法

核心代码实现 #

以下是完全符合需求的自定义Dataset类，适配Google Colab环境：

import os
import scipy.io
import torch
from torch.utils.data import Dataset

class SPADDataset(Dataset):
    def __init__(self, data_dir):
        # 遍历目标文件夹，获取所有.mat文件路径
        self.mat_files = [os.path.join(data_dir, f) for f in os.listdir(data_dir) 
                          if f.endswith('.mat')]
        # 提前读取所有数据并转换为Tensor（按需求在__init__中完成）
        self.spad_data = []
        self.depth_data = []
        
        for mat_path in self.mat_files:
            # 读取.mat文件
            mat_dict = scipy.io.loadmat(mat_path)
            # 提取SPAD直方图和深度图数据（键名需与.mat文件内变量名一致）
            spad_np = mat_dict['SPAD']
            depth_np = mat_dict['depth']
            
            # 转换为PyTorch Tensor，统一数据类型为float32
            spad_tensor = torch.from_numpy(spad_np).float()
            depth_tensor = torch.from_numpy(depth_np).float()
            
            self.spad_data.append(spad_tensor)
            self.depth_data.append(depth_tensor)

    def __len__(self):
        # 返回数据集总长度
        return len(self.mat_files)

    def __getitem__(self, idx):
        # 返回指定索引的SPAD数据和深度图
        return self.spad_data[idx], self.depth_data[idx]

关键细节说明 #

• 遍历文件夹：通过os.listdir筛选.mat后缀文件，用os.path.join拼接完整路径，避免跨平台路径错误。
• 读取.mat文件：scipy.io.loadmat返回字典结构，直接通过键名（SPAD、depth）提取对应numpy数组，注意键名必须与.mat文件内存储的变量名完全匹配。
• Tensor转换：用torch.from_numpy()将numpy数组转为Tensor，统一转为float32类型，适配PyTorch模型训练的常规数据类型要求。
• Dataset接口实现：__len__返回数据总量，__getitem__返回单样本的SPAD张量和深度张量，直接支持DataLoader迭代调用。

Google Colab使用示例 #

若数据存储在Google Drive中，先挂载Drive再创建Dataset和DataLoader：

# 挂载Google Drive（仅首次运行需要）
from google.colab import drive
drive.mount('/content/drive')

# 指定数据集文件夹路径（根据实际路径修改）
data_dir = '/content/drive/MyDrive/SPAD_Counts'

# 创建Dataset实例
dataset = SPADDataset(data_dir)

# 创建DataLoader用于训练迭代
from torch.utils.data import DataLoader
dataloader = DataLoader(dataset, batch_size=8, shuffle=True)

# 测试迭代逻辑
for spad_batch, depth_batch in dataloader:
    print(f"SPAD batch shape: {spad_batch.shape}")  # 输出应为[8,64,64,128]
    print(f"Depth batch shape: {depth_batch.shape}") # 输出应为[8,64,64]
    break

优化建议（可选） #

若数据集文件数量多、单文件体积大，提前在__init__加载所有数据会占用大量内存，建议改为延迟加载（仅在__getitem__时读取数据），修改后的核心代码如下：

def __init__(self, data_dir):
    self.mat_files = [os.path.join(data_dir, f) for f in os.listdir(data_dir) 
                      if f.endswith('.mat')]

def __getitem__(self, idx):
    mat_path = self.mat_files[idx]
    mat_dict = scipy.io.loadmat(mat_path)
    spad_tensor = torch.from_numpy(mat_dict['SPAD']).float()
    depth_tensor = torch.from_numpy(mat_dict['depth']).float()
    return spad_tensor, depth_tensor

这种方式仅在需要时读取数据，大幅降低内存占用，适合大规模数据集场景。

内容的提问来源于stack exchange，提问作者Anand Idris