这篇文章是讲述从mmdetection的dataLoader所生成的data字典(包含DataContainer对象)里面获取tensor以及shape的方法

需求:

如果我的batchsize是8,而我有50个数据,那么前6个batch都包含8个数据,最后一个batch有2个.

我需要每个batch都进行一次FPS计算，就要知道一个batch有多少个数据.

困难点:

mmdetection dataloader 产生的数据对象是一个字典, 且包含DataContainer类别对象，不能直接看到数据的shape是多少

要精准知道一个DataContainer包含多少个数据。就需要知道如何从DataContainer对象获取tensor类型数据, 然后才能获取shape, 才能获取batchsize

（以下代码均为mmdetection的代码，地址：https://github.com/open-mmlab/mmdetection/blob/master/tools/analysis_tools/benchmark.py
点击跳转）

其中，这是dataloader的定义

data_loader = build_dataloader(dataset,# samples_per_gpu=1,# 这里我修改了samples_per_gpu=samples_per_gpu,# Because multiple processes will occupy additional CPU resources,# FPS statistics will be more unstable when workers_per_gpu is not 0.# It is reasonable to set workers_per_gpu to 0.workers_per_gpu=0,dist=True,shuffle=False)

这是dataloader用for循环产生data(一个字典, 其中的DataContainer对象才包含数据tensor)来做inference的方法

    # benchmark with 2000 image and take the averagefor i, data in enumerate(data_loader):torch.cuda.synchronize()start_time = time.perf_counter()with torch.no_grad():model(return_loss=False, rescale=True, **data)torch.cuda.synchronize()elapsed = time.perf_counter() - start_time# print(i, f"data shape {data['img'][0].data[0].shape}")# 完成这个batch的infer之后会算一次时间if i >= num_warmup:pure_inf_time += elapsedif (i + 1) % log_interval == 0:# 这里我也修改了# fps = (i + 1 - num_warmup) / pure_inf_timefps = (i + 1 - num_warmup) / (pure_inf_time / samples_per_gpu)print(f'Done image [{i + 1:<3}/ {max_iter}], 'f'fps: {fps:.1f} img / s, 'f'times per image: {1000 / fps:.1f} ms / img',flush=True)if (i + 1) == max_iter:fps = (i + 1 - num_warmup) / pure_inf_timeprint(f'Overall fps: {fps:.1f} img / s, 'f'times per image: {1000 / fps:.1f} ms / img',flush=True)breakreturn fps

不想看详细过程的朋友,我这里给出结论:

这个for循环里面,使用

print(data[‘img’][0].data[0].shape)

就可以看到shape了, 以下是结果

torch.Size([16, 3, 608, 608])

接下来继续讲述过程

在这个for循环下,print(data)会发现data是一个字典打印出一大堆

{'img_metas': [DataContainer([[{'filename': 'dat .........

然后用print(data.key())来查看有哪些键,发现打印出

dict_keys(['img_metas', 'img'])

用print(data[‘img’]) 来查看内容(终于看到tensor了…)

[DataContainer([tensor([[[[0.4667, 0.4745, 0.4980,  ..., 0.1373, 0.1255, 0.1255],[0.4627, 0.4588, 0.5098,  ..., 0.1647, 0.1804, 0.1647],[0.4784, 0.4941, 0.4706,  ..., 0.1843, 0.1686, 0.1804],...,[1.0000, 1.0000, 1.0000,  ..., 0.4157, 0.4431, 0.4353],[1.0000, 1.0000, 1.0000,  ...

发现还被一个list包裹着, 那就使用 data[‘img’][0], 这个就是DataContainer对象了,查文档发现DataContainer对象有一个data属性,我猜这个就能够提取出tensor了

    @propertydef data(self) -> Union[torch.Tensor, np.ndarray]:return self._data

使用 print(data[‘img’][0].data),发现还是被一个list包裹,根据上面的资料,这个data属性会返回一个[torch.Tensor, np.ndarray],包含两个元素的列表,那么用data[‘img’][0].data[0]就可以拿到tensor了

[tensor([[[[0.4667, 0.4745, 0.4980,  ..., 0.1373, 0.1255, 0.1255],[0.4627, 0.4588, 0.5098,  ..., 0.1647, 0.1804, 0.1647],[0.4784, 0.4941, 0.4706,  ..., 0.1843, 0.1686, 0.1804],...,[1.0000, 1.0000, 1.0000,  ..., 0.

那么print(data[‘img’][0].data[0].shape),发现可以了

torch.Size([16, 3, 608, 608])