【DETR系列目标检测算法代码精讲】01 DETR算法03 Dataloader代码精讲

news2024/11/25 22:57:53

与一般的Dataloader的区别在于我们对图像进行了随机裁剪，需要进行额外的操作才能将其打包到dataloader里面

在这里插入图片描述
这一段的代码如下：

    if args.distributed:
        sampler_train = DistributedSampler(dataset_train)
        sampler_val = DistributedSampler(dataset_val, shuffle=False)
    else:
        sampler_train = torch.utils.data.RandomSampler(dataset_train)
        sampler_val = torch.utils.data.SequentialSampler(dataset_val)

    batch_sampler_train = torch.utils.data.BatchSampler(sampler_train, args.batch_size, drop_last=True)

    data_loader_train = DataLoader(dataset_train,
                                   batch_sampler=batch_sampler_train,
                                   collate_fn=utils.collate_fn,
                                   # num_workers=args.num_workers
                                   )
    data_loader_val = DataLoader(dataset_val,
                                 args.batch_size,
                                 sampler=sampler_val,
                                 drop_last=False,
                                 collate_fn=utils.collate_fn,
                                 # num_workers=args.num_workers
                                 )

对于训练数据集，使用RandowSampler类进行随机采样
对于验证数据集，使用SequentialSampler进行顺序采样

采样以后，使用BatchSampler打包成batch
然后再使用Dataloader

Dataloader中有个函数
collate_fn

在这里插入图片描述
这个函数又调用了一个函数nested_tensor_from_tensor_list
这个函数重新定义了我们输入数据的格式

在这里插入图片描述

默认的batch为2，我们输入的就是包含了两个元素的list，其中每个元素都是我们从dataset的__getitem__方法获得的输出

然后通过zip函数进行解析
在这里插入图片描述
可以看到之前的形式是一个元素中是img+target
现在变成了一个元素里面都是img，另一个元素里面都是target
batch[0] = nested_tensor_from_tensor_list(batch[0])
然后通过索引0取出图像部分

传入到nested_tensor_from_tensor_list方法中
这个方法的全部代码如下

def nested_tensor_from_tensor_list(tensor_list: List[Tensor]):
    # TODO make this more general
    if tensor_list[0].ndim == 3:
        if torchvision._is_tracing():
            # nested_tensor_from_tensor_list() does not export well to ONNX
            # call _onnx_nested_tensor_from_tensor_list() instead
            return _onnx_nested_tensor_from_tensor_list(tensor_list)

        # TODO make it support different-sized images
        max_size = _max_by_axis([list(img.shape) for img in tensor_list])

        batch_shape = [len(tensor_list)] + max_size
        b, c, h, w = batch_shape
        dtype = tensor_list[0].dtype
        device = tensor_list[0].device
        tensor = torch.zeros(batch_shape, dtype=dtype, device=device)
        mask = torch.ones((b, h, w), dtype=torch.bool, device=device)
        for img, pad_img, m in zip(tensor_list, tensor, mask):
            pad_img[: img.shape[0], : img.shape[1], : img.shape[2]].copy_(img)
            m[: img.shape[1], :img.shape[2]] = False
    else:
        raise ValueError('not supported')
    return NestedTensor(tensor, mask)

定义了一个函数nested_tensor_from_tensor_list，该函数接受一个Tensor列表（tensor_list）作为输入，并返回一个NestedTensor对象。NestedTensor是一个特殊的数据结构，通常用于表示图像数据，其中可以包含不同大小的图像，并且有一个与之对应的掩码（mask）来表示每个图像的实际大小。

检查tensor_list中的第一个Tensor的维度是否为3。

通过_max_by_axis函数来确定tensor_list中所有图像的最大尺寸。这意味着最终生成的NestedTensor将包含所有图像的最大高度和宽度。

这个_max_by_axis函数的代码如下：

def _max_by_axis(the_list):
    # type: (List[List[int]]) -> List[int]
    maxes = the_list[0]
    for sublist in the_list[1:]:
        for index, item in enumerate(sublist):
            maxes[index] = max(maxes[index], item)
    return maxes