从mmdetection的dataLoader所生成的data字典(包含DataContainer对象)里面获取tensor以及shape的方法

这篇文章是讲述从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 average
    for 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 += elapsed
            if (i + 1) % log_interval == 0:
            	# 这里我也修改了
            	# fps = (i + 1 - num_warmup) / pure_inf_time
                fps = (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_time
            print(
                f'Overall fps: {fps:.1f} img / s, '
                f'times per image: {1000 / fps:.1f} ms / img',
                flush=True)
            break
    return 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了

    @property
    def 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])