Tensor

常用属性

1
2
3
4
5
6
tensor.type()    # 张量的数据类型
tensor.size()    # 张量的形状,返回类型是 tuple 的子类
tensor.dim()     # 张量的维度

# 设置默认张量类型,主要 PyTorch 里的 float 远远快于 double
torch.set_default_tensor_type(torch.FloatTensor)

数据转换

CPU to GPU

1
2
tensor = tensor.cuda()  # To device: GPU
tensor = tensor.cpu()   # To device: CPU

注意:

  • 对于 Tensor 来说,返回的是该 Tensor 在显存上的拷贝,而不改变自身;
  • 对于 nn.Module 来说,调用 cuda() 是将自身放入显存,此时该操作是 in-place 的。

Tensor to Numpy

1
2
3
4
# Tensor to ndarray
ndarray = tensor.cpu().numpy()
# ndarray to Tensor
tensor = torch.from_numpy(ndarray).float()

Tensor to PIL

可以简单使用 torchvision 库:

1
2
3
4
5
# Tensor to Image
image = torchvision.transforms.functional.to_pil_image(tensor)
tensor =
# Image to Tensor
torchvision.transforms.functional.to_tensor(PIL.Image.open(path))

也可以自己转,不过要注意两者的数据维度和数值范围,自己做维度调整与归一化:

  • Tensor 默认是 (N))CHW,数据范围 [0,1]
  • PIL.Image 默认是 HWC,数据范围 [0, 255]
1
2
3
4
5
6
7
8
# Tensor to Image
image = PIL.Image.fromarray(
    torch.clamp(tensor * 255, min=0, max=255).byte().permute(1, 2, 0).cpu().numpy()
    )
# Image to Tensor
tensor = torch.from_numpy(
    np.asarray(PIL.Image.open(path))
    ).permute(2, 0, 1).float() / 255

Numpy to PIL

1
2
3
4
# ndarray to Image
image = PIL.Image.fromarray(ndarray.astypde(np.uint8))
# Image to ndarray
ndarray = np.asarray(PIL.Image.open(path))

常用操作

拷贝/取出计算图

1
2
3
4
tensor.clone()   # 张量的拷贝(放在新的内存),依然存在于计算图中
tensor.detach()  # 张量从计算图中取出(内存与原张量是共享的)

tensor.item()  # 以 Python 数据类型返回一个单值张量(取出了计算图),常用于 loss、accuracy 等

拼接

1
2
tensor = torch.cat(list_of_tensors, dim=0)    # 拼接张量,不增加新维度
tensor = torch.stack(list_of_tensors, dim=0)  # 拼接张量,增加新维度

水平翻转

1
2
# 假设 tensor 的形状为 NCHW
tensor = tensor[:, :, :, torch.arange(tensor.size(3) - 1, -1, -1).long()]

注意 PyTorch 不支持 tensor[::-1] 这样的负步长操作,因此水平翻转可以用张量索引实现。

张量扩展

1
2
torch.allclose(tensor1, tensor2)  # float tensor
torch.equal(tensor1, tensor2)     # int tensor

得到零/非零元素

1
2
3
4
torch.nonzero(tensor)               # Index of non-zero elements
torch.nonzero(tensor == 0)          # Index of zero elements
torch.nonzero(tensor).size(0)       # Number of non-zero elements
torch.nonzero(tensor == 0).size(0)  # Number of zero elements

判断张量相等

1
2
torch.allclose(tensor1, tensor2)  # float tensor
torch.equal(tensor1, tensor2)     # int tensor

One-hot 转换

1
2
3
4
5
6
7
N = tensor.size(0)
one_hot = torch.zeros(N, num_classes).long()
one_hot.scatter_(
    dim=1,
    index=torch.unsqueeze(tensor, dim=1),
    src=torch.ones(N, num_classes).long()
)

欧氏距离

1
2
# X1 is of shape m*d, X2 is of shape n*d.
dist = torch.sqrt(torch.sum((X1[:,None,:] - X2) ** 2, dim=2))

DataLoader

lmdb, TFRecord

PyTorch 自带的 DataLoader 效率较低,可以结 Caffe 的 lmdb、TensorFlow 的 TFRecord 等数据存储与读取。

pin_memory

  • 当内存充足时,DataLoader 尽量使用 pin_memory=True,开启后通过DataLoader 生成的 Tensor 数据会放入主机中的 锁页内存(即任何情况下都不会放入虚拟内存),避免速度被拖慢。这里说的「内存」指的是主机内存,显卡内存可以理解为都是锁页内存。该配置默认为 False。

__len__()

DataLoader 有时会间歇式出现卡顿,可以手动设置其长度 __len__(self) 避免:

1
2
def __len__(self):
        return self.images.shape[0]

其他细节

  • 第一个 GPU 占用的显存总是较多一点,可以使用开源工具 PyTorch-Encoding 来平衡。
  • 可以定期使用 torch.cuda.empty_cache()del <variable> 清显存。

Flags

cudnn.enabled

torch.backends.cudnn.enabled = True

允许使用 cuDNN。

cudnn.benchmark

1
torch.backends.cudnn.benchmark = True

启用 cuDNN 在首次执行时 自动优化计算图,从而加速计算速度,一般可达 10~20%,视具体情况而定,训练与预测都可以加速。

具体来说,优化的计算图是 卷积层的实现算法,因为卷积层的效率与自身参数、输入数据、硬件平台、执行精度、显存内数据布局、具体执行卷积的方式等等因素有关,而设定该 flag 之后 cuDNN 会自动地根据当前配置尝试不同的卷积实现算法,从而得到最高效的卷积实现。

但是该步骤需要的时间较长,一般在 网络结构固定(非动态改变)、网络的输入/中间层/输出大小不变 时使用,否则每次 cuDNN 都会重新优化,反而导致运行更慢,甚至内存不足。

该 flag 默认为 False。

更详细内容可以参考:

cudnn.deterministic

1
torch.backends.cudnn.deterministic = False

cuDNN 的计算过程中有随机性,设置该 flag 可以固定随机结果,便于复现等。

代码片段

查看 CUDA 版本

1
2
3
4
torch.__version__               # PyTorch 版本
torch.version.cuda              # CUDA 版本
torch.backends.cudnn.version()  # cuDNN 版本
torch.cuda.get_device_name(0)   # GPU 类型

可复现实验设置

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
# random
random.seed(0)

# Numpy
np.random.seed(0)
np.random.RandomState(0)  # If using RandomState

# PyTorch
torch.manual_seed(0)  # For both CPU and CUDA
torch.backends.cudnn.benchmark = False  # If using CUDNN
torch.backends.cudnn.deterministic = True  # If using CUDNN

读写 h5 数据

读取 h5:

1
2
3
4
5
with h5py.File(outfile) as f:
    imagename = [x.decode() for x in f['imagename']]
    kp2ds = np.array(f['part'])
    kp3ds = np.array(f['S'])
    cvimgs = np.array(f['image'])

写入 h5:

1
2
3
4
5
with h5py.File(outfile) as f:
    imagename = [x.decode() for x in f['imagename']]
    kp2ds = np.array(f['part'])
    kp3ds = np.array(f['S'])
    cvimgs = np.array(f['image'])

常见错误

Insufficient shm

1
ERROR: Unexpected bus error encountered in worker. This might be caused by insufficient shared memory (shm)

原因:

在 Docker 中运行代码时,由于 batch size 设置得过大,shared memory 不够(Docker 限制了 shm)。

解决方式

Dataloadernum_workers 设置为 0。

参考