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

本文主要是介绍【DETR系列目标检测算法代码精讲】01 DETR算法03 Dataloader代码精讲，希望对大家解决编程问题提供一定的参考价值，需要的开发者们随着小编来一起学习吧！

与一般的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 generalif 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() insteadreturn _onnx_nested_tensor_from_tensor_list(tensor_list)# TODO make it support different-sized imagesmax_size = _max_by_axis([list(img.shape) for img in tensor_list])batch_shape = [len(tensor_list)] + max_sizeb, c, h, w = batch_shapedtype = tensor_list[0].dtypedevice = tensor_list[0].devicetensor = 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]] = Falseelse: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