Pytorch常用API解析

发表于 | 分类于

好记性不如烂笔头，主要是纪录一些Pytorch中常用函数用法以及自定义数据读取collate_fn()介绍。

常用数学类操作符

torchTensor.transpose() & Tensor.permute()

transpose只能操作矩阵的两个维度。只接受两个维度参数。若要对多个维度进行操作，使用permute更加灵活方便。

1
2
3
4
5
a=torch.rand(2,3,4)  #torch.Size([2, 3, 4])
b=a.transpose(0,1)   #torch.Size([3, 2, 4])

c=a.transpose(0,1).transpose(1,2)  #torch.Size([3, 4, 2])
d=a.permute(1,2,0)                 #torch.Size([3, 4, 2])

同样是对三个维度进行变换，transpose需要操作三次，而是用permute只需要操作一次，因此对于高维度的矩阵变化，permute更加方便。

toch.cat() & torch.stack()

cat是对数据沿着某一维度进行拼接。cat后数据的总维数不变。二维的矩阵拼接后依旧是二维的矩阵。如下面代码对两个2维tensor（分别为23,43）进行拼接，拼接完后变为3*3还是2维的tensor。

1
2
3
x=torch.rand(2,3)
y=torch.rand(4,3)
z=torch.cat((x,y),0)  #torch.Size([6, 3])

注意：对维度进行拼接是必须满足维度一致的要求，除了指定拼接的维度除外，其他维度必须都相等。否则就会出现维度不匹配问题。
stack是增加新的维度进行堆叠。这个比较常见，比如在pytorch中对数据进行加载组成一个batch时常用到。比较好理解。

1
2
3
4
a=torch.rand(3,224,224)
b=torch.rand(3,224,224)
c=torch.stack((a,b),0)  # torch.Size([2, 3, 224, 224])
d=torch.stack((a,b),3)  # torch.Size([3, 224, 224, 2])

注意cat()和stack()两者区别就在于后者会增加维度。

torch.squeeze() & torch.unsqueeze()

squeeze(dim_n)压缩，即去掉元素数量为1的dim_n维度。同理unsqueeze(dim_n)，增加dim_n维度，元素数量为1。

1
2
3
4
5
6
7
8
#减少维度
a=torch.rand(2,1,4)  #torch.Size([2, 1, 4])
a_=a.squeeze()       #torch.Size([2, 4])，不加参数，去掉所有为元素个数为1的维度
a_=a.squeeze(0)      #torch.Size([2, 1, 4])加上参数，去掉第一维的元素为1，不起作用，因为第一维有2个元素
a_=a.squeeze(1)      #torch.Size([2, 4])，这样就可以

#增加维度
a_=a.unsqueeze(0)    #torch.Size([1, 2, 1, 4])

Tensor.expand_as()

expand_as()这是tensor变量的一个内置方法，如果使用b.expand_as(a)就是将b进行扩充，扩充到a的维度，需要说明的是a的低维度需要比b大，例如b的shape是31，如果a的shape是32不会出错，但是是2*2就会报错了

1
2
3
4
5
6
a=torch.Tensor([[1],[2],[3]])
b=torch.Tensor([[4]])
c=b.expand_as(a)
#tensor([[4.],
        [4.],
        [4.]])

torch.contiguous()

contiguous：view只能用在contiguous的variable上。如果在view之前用了transpose, permute等，需要用contiguous()来返回一个contiguous copy。 因为view需要tensor的内存是整块的。有些tensor并不是占用一整块内存，而是由不同的数据块组成，而tensor的view()操作依赖于内存是整块的，这时只需要执行contiguous()这个函数，把tensor变成在内存中连续分布的形式。判断是否contiguous用torch.Tensor.is_contiguous()函数

1
2
3
4
x=torch.ones(10,10)
x.is_contiguous()  #True
x.transpose(0,1).is_contiguous()  #False
x.transpose(0, 1).contiguous().is_contiguous()  #True

因此，在调用view之前最好先contiguous一下，x.contiguous().view() 。

数据读取

在pytorch中一个现有的数据读取方法就是torchvision.datasets.ImageFolder()，这个api主要用于做分类问题。将每一类数据放到同一个文件夹中，比如有10个类别，那么就在一个大的文件夹下面建立10个子文件夹，每个子文件夹里面放的是同一类的数据。从api的实现可以发现ImageFolder该类继承自torch.utils.data.Dataset。

Dataset

Dataset的定义如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
class Dataset(object):
    """An abstract class representing a Dataset.

    All other datasets should subclass it. All subclasses should override
    ``__len__``, that provides the size of the dataset, and ``__getitem__``,
    supporting integer indexing in range from 0 to len(self) exclusive.
    """

    def __getitem__(self, index):
        raise NotImplementedError

    def __len__(self):
        raise NotImplementedError

    def __add__(self, other):
        return ConcatDataset([self, other])

从上面的注释知道，这个是代表数据集的一个抽象类。有关于数据集的类都可以定义为其子类，只需要重写getitem和len就可以了。
那么定义好了数据集我们不可能将所有的数据集都放到内存，这样内存肯定就爆了，我们需要定义一个迭代器，每一步产生一个batch，这里PyTorch已经为我们实现好了，就是下面的torch.utils.data.DataLoader。

DataLoader

DataLoader能够为我们自动生成一个多线程的迭代器。DataLoader的函数定义如下：DataLoader(dataset, batch_size=1, shuffle=False, sampler=None, num_workers=0, collate_fn=default_collate, pin_memory=False, drop_last=False)

  • dataset：加载的数据集(Dataset对象)
  • batch_size：batch size
  • shuffle:：是否将数据打乱
  • sampler： 样本抽样，后续会详细介绍
  • num_workers：使用多进程加载的进程数，0代表不使用多进程
  • collate_fn： 如何将多个样本数据拼接成一个batch，一般使用默认的拼接方式即可
  • pin_memory：是否将数据保存在pin memory区，pin memory中的数据转到GPU会快一些
  • drop_last：dataset中的数据个数可能不是batch_size的整数倍，drop_last为True会将多出来不足一个batch的数据丢弃

collate_fn()

到这里，我们可以发现pytorch数据载入很简单，基本上对于简单地任务都不会出啥问题，然而数据载入容易出错的地方往往在于collate_fn函数。尤其是对于batch大小不一致的情况，比如目标检测，需要检测出图片中所有的目标。如下：

问题就是，输入的是一张张图片，大小不同可以resize成一样大小组成一个batch，但是它的label是每个目标的boxes和categories。而且每张图片中的目标个数也不一样，怎么组成一个batch。这个时候我们就需要自定义实现一个collate_fn函数。这里可以使用任何名字，只要在DataLoader里面传入就可以了。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
def collate_fn(self, batch):
    """
    Since each image may have a different number of objects, we need a collate function (to be passed to the DataLoader).

    This describes how to combine these tensors of different sizes. We use lists.

    Note: this need not be defined in this Class, can be standalone.

    :param batch: an iterable of N sets from __getitem__()
    :return: a tensor of images, lists of varying-size tensors of bounding boxes, labels, and difficulties
    """
    
    images = list()
    boxes = list()
    labels = list()
    difficulties = list()

    for b in batch:
        images.append(b[0])
        boxes.append(b[1])
        labels.append(b[2])
        difficulties.append(b[3])

    images = torch.stack(images, dim=0)

    return images, boxes, labels, difficulties  # tensor (N, 3, 300, 300), 3 lists of N tensors each

pin_memory()

上面提到了pin_memory参数，pin_memory就是锁页内存，创建DataLoader时，设置pin_memory=True，则意味着生成的Tensor数据最开始是属于内存中的锁页内存，这样将内存的Tensor转义到GPU的显存就会更快一些。

主机中的内存，有两种存在方式，一是锁页，二是不锁页，锁页内存存放的内容在任何情况下都不会与主机的虚拟内存进行交换（注：虚拟内存就是硬盘），而不锁页内存在主机内存不足时，数据会存放在虚拟内存中。

而显卡中的显存全部是锁页内存！

当计算机的内存充足的时候，可以设置pin_memory=True。当系统卡住，或者交换内存使用过多的时候，设置pin_memory=False。因为pin_memory与电脑硬件性能有关，pytorch开发者不能确保每一个炼丹玩家都有高端设备，因此pin_memory默认为False。