mmdet3d读入waymo dataset:from file to input tensor

本文介绍了如何使用MMDet3D从Waymo数据集读取数据并转换为输入张量。内容涵盖了torch.utils.data模块,包括迭代器、数据集(如Waymo Dataset)、Sampler和DataLoader的工作原理。着重讨论了数据从文件系统到GPU的处理流程,包括数据加载、预处理、增强和模型训练过程。

摘要生成于 C知道 ,由 DeepSeek-R1 满血版支持, 前往体验 >

之前已经处理出了kitti format的waymo dataset文件,但距离模型真正处理数据还有很长一段路。虽然在runner的train里,data的获取很简单,就是一个enumerate(dataloader):

for i, data_batch in enumerate(self.data_loader):
            self._inner_iter = i
            self.call_hook('before_train_iter')
            self.run_iter(data_batch, train_mode=True, **kwargs)
            # model forward && calc loss && backprop
            self.call_hook('after_train_iter')
            # optimize, save logs
            self._iter += 1

但事实上数据要从file system读入到内存,经过pipeline类的格式处理、数据增强,进入gpu,然后才能给已经load进gpu的model进行训练。
在这个过程中有几个很关键的类在起作用,有dataset类,sampler类,dataloader类,mmdet的这几个类基本上继承自pytorch类。接下来结合waymo的代码介绍一下这三个类。

torch.utils.data

iterator 迭代器

和c++里的set等stl类一样,iterator可以帮助我们遍历整个类,或者说容器的内容。在这里主要是sampler和dataloader有自己的iterator,这样就能实现for i,data in enumerate(dataloader): output = model(data)的功能了。
对于dataloader来说,为了支持iterator,需要实现几个方法:

  1. .__len__(self),一般返回类的长度,比如数据有几个frame。
  2. .__getitem__(self),定义获取容器指定元素的行为,这样就可以以data[index]的下标方式访问类里的信息。
  3. .__iter__(self):定义当迭代容器中的元素时的行为,我的理解是,这样使得for之初能获取到dataloader的iterator
  4. .__next__(self):定义迭代器的迭代规则,比如按某种顺序遍历整个set,for循环每次都会调用一下next。
    具体例子可以看教程一

dataset

主要实现了数据从文件读到内存的功能,定义了__getitem__,这样当你dataset[index]的时候,就能得到某个frame的数据。主要类型有两种&

首先,需要将数据读入并进行归一化处理,可以使用以下代码: ```python import os import numpy as np def read_data(folder_path): data = [] for filename in os.listdir(folder_path): file_path = os.path.join(folder_path, filename) arr = np.loadtxt(file_path) arr = arr / 255.0 # 归一化 data.append(arr) return np.array(data) input_folder = "input_folder" output_folder = "output_folder" input_data = read_data(input_folder) output_data = read_data(output_folder) ``` 接下来,可以使用PyTorch实现Densenet模型。以下是一个简单的Densenet实现: ```python import torch import torch.nn as nn import torch.nn.functional as F from torchvision.models import densenet121 class Densenet(nn.Module): def __init__(self): super(Densenet, self).__init__() self.densenet = densenet121(pretrained=True) self.linear = nn.Linear(1000, 4096) def forward(self, x): x = self.densenet(x) x = self.linear(x) return x ``` 接下来,可以定义数据集和数据加载器,使用PyTorch的内置函数进行训练。以下是一个简单的训练过程: ```python from torch.utils.data import Dataset, DataLoader class MyDataset(Dataset): def __init__(self, input_data, output_data): self.input_data = input_data self.output_data = output_data def __len__(self): return len(self.input_data) def __getitem__(self, idx): input_arr = self.input_data[idx] output_arr = self.output_data[idx] return input_arr, output_arr train_ratio = 0.8 train_size = int(len(input_data) * train_ratio) train_input = input_data[:train_size] train_output = output_data[:train_size] test_input = input_data[train_size:] test_output = output_data[train_size:] train_dataset = MyDataset(train_input, train_output) test_dataset = MyDataset(test_input, test_output) batch_size = 10 train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True) test_loader = DataLoader(test_dataset, batch_size=batch_size) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = Densenet().to(device) optimizer = torch.optim.Adam(model.parameters(), lr=0.001) criterion = nn.MSELoss() num_epochs = 10 for epoch in range(num_epochs): train_loss = 0.0 for batch_idx, (input_arr, output_arr) in enumerate(train_loader): input_arr = input_arr.to(device) output_arr = output_arr.to(device) optimizer.zero_grad() output = model(input_arr.unsqueeze(1).float()) loss = criterion(output, output_arr.unsqueeze(1).float()) loss.backward() optimizer.step() train_loss += loss.item() train_loss /= len(train_loader) test_loss = 0.0 with torch.no_grad(): for batch_idx, (input_arr, output_arr) in enumerate(test_loader): input_arr = input_arr.to(device) output_arr = output_arr.to(device) output = model(input_arr.unsqueeze(1).float()) loss = criterion(output, output_arr.unsqueeze(1).float()) test_loss += loss.item() test_loss /= len(test_loader) print("Epoch {} Train Loss {:.6f} Test Loss {:.6f}".format(epoch+1, train_loss, test_loss)) ``` 最后,可以实现一个函数,输入一个二维数组,输出映射后的二维数组: ```python def map_array(arr): arr = arr / 255.0 # 归一化 arr = torch.tensor(arr).unsqueeze(0).unsqueeze(0).float().to(device) with torch.no_grad(): output = model(arr) return output.squeeze(0).squeeze(0).cpu().numpy() * 255.0 ``` 这样,就可以使用以上代码实现灰度图像到灰度图像的Densenet映射了。
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值