这里简单介绍一下用PyTorch在CPU上的一些性能相关的BKM。
内容以inference为主，毕竟CPU上主要的场景还是inference；另外这里CPU都指的是Intel Xeon.
gist里面写了英文版的，内容和这里的基本相当： General guidelines for CPU performance on PyTorch
1. 使用mkldnn layout
layout是指数据在tensor中的组织形式，PyTorch默认的layout是NCHW，也就是{batch_size, channel, height, width}。MKL-DNN（intel在CPU上针对Neural Network相关operator的加速库，最近改名叫DNNL了，这里还是叫MKL-DNN吧，毕竟叫这么久了。。）从performance优化角度设计，对于不同的convolution pattern和algorithm会选择不同的layout进行计算，就是所谓的mkldnn layout，有时也叫internal layout或者primitive layout。与其相对应的PyTorch中默认的layout叫default layout或者native layout。default layout和mkldnn layout之间的转换叫做reorder。其实通过layout reorder的方式加速convolution是常规手段，比如NHWC天生就比NCHW快。但是MKL-DNN搞得比较复杂，有很多internal的layout，比如nChw16c，指的是把4D的NCHW转成5D的，把C这个维度除16，这样最后一维比较容易做vectorization，AVX512有512-bit的指令位宽，刚好是16个float。
不过MKL-DNN的layout搞得这么复杂，有个麻烦：用户希望看到的永远都是default的layout。如果输出一个mkldnn internal layout其实对用户没有意义，这就需要把mkldnn layout“藏”起来。这个“藏”不是个轻松地工作，static graph可以通过改图完成，PyTorch是dymanic graph，没有所谓的图给你改，所以是通过dynamic dispatch来做的，mkldnn layout相当于dispatch的路标，这个概念实际上叫做TensorTypeId，这里不详细展开。
PyTorch CPU上面默认情况下Conv2d也是使用MKL-DNN来跑，不过每次需要把input和weight从default转成mkldnn的layout，计算完的output再从mkldnn转成default layout。这里面是有很大的性能损失的，性能大概要减半。
为了达到更好的performance，就需要“允许”mkldnn layout在不同的operator之间流转，简而言之就是对于MKL-DNN支持的operator，使用mkldnn layout，对于其不支持的operator转回default layout。主要涉及两个方面：
- 把input转成mkldnn layout，为了消除operator之间的default2mkldnn和mkldnn2default转换开销；
2. 把model转成mkldnn版本，为了inference时候做weight cache做准备，模型转换时会将operator换成MKL-DNN版本，e.g. Conv2d -> MkldnnConv2d；还会将weight转成mkldnn layout.
这里有个问题是MKL-DNN只支持二十来个operator，比如像Conv2d, BatchNorm，ReLU这些；而PyTorch里面有上百个operator。怎么办呢？
如果model里面用到的operator都是支持mkldnn layout的，那么很简单：
from torch.utils import mkldnn as mkldnn_utils
# change input to be mkldnn layout
input_ = input.to_mkldnn()
# change weight to be mkldnn layout
model_ = mkldnn_utils.to_mkldnn(model)
# NB: output_ is mkldnn layout
output_ = model_(input_)
# change output_ back to default layout
output = output_.to_dense()整个过程和用cuda的流程类似，只不过是把cuda(), cpu()换成了to_mkldnn()以及to_dense()。
如果model里面有的operator支持mkldnn layout而有的不支持，那就麻烦一些，需要手动在forward里面插入to_mkldnn()和to_dense()：
class MyModel(nn.Module):
def __init__(self):
self(MyModel, self).__init__()
self.conv1 = nn.Conv2d(10, 10, 3)
# MyModel has mkldnn unsupported operators X()
self.unsupported_mod = nn.X()
self.linear1 = nn.Linear(10, 20)
def forward(self, x):
x = self.conv1(x)
# use default layout for module without mkldnn support
x = x.to_dense()
x = self.unsupported_mod(x)
x = x.to_mkldnn()
x = self.linear1(x)
return x这个过程你可以想象成cuda的backend有个operator不支持，需要在cpu上面跑，处理方式是类似的。
Notes:
- to_mkldnn()和to_dense()是有时间开销的，会影响performance，这是memory copy。
- mkldnn layout不支持view()，需要view()的地方要换成reshape()，同样这个reshape也是有时间开销的。
2. 例子：deploy resnext101
下面这个拿torchvision里面的resnext101_32x8d 举例说明如果完成CPU上的部署，这个是FB要求的重点模型。 code在这里：convnet-benchmark-py，用的batch_size为1。
2.a default:
./run.sh --inference --single结果是 92ms/image:
ModelType: resnext101, Kernels: nn Input shape: 1x3x224x224
nn :forward: 92.52 (ms) 10.81 (imgs/s)
nn :backward: 0.00 (ms)
nn :update: 0.00 (ms)
nn :total: 92.52 (ms) 10.81 (imgs/s)2.b 使用mkldnn layout
./run.sh --inference --single --mkldnn这里会完成input以及model的转换：
from torch.utils import mkldnn as mkldnn_utils
input = input.to_mkldnn() # input will be _mkldnn layout
model = mkldnn_utils.to_mkldnn(model) # weight will be _mkldnn layout结果是 43ms/image:
ModelType: resnext101, Kernels: nn Input shape: 1x3x224x224
nn :forward: 43.27 (ms) 23.11 (imgs/s)
nn :backward: 0.00 (ms)
nn :update: 0.00 (ms)
nn :total: 43.27 (ms) 23.11 (imgs/s)2.c 使用weight cache
./run.sh --inference --single --mkldnn --cache-weight这里会将model转化成TorchScript，接下来save成一个 .pt 文件同时完成对weight的cache，再load进来的model就不会再做weight reorder。
traced = torch.jit.trace(net, data, check_trace=False)
script = traced.save('model.pt') # mkldnn reordered weight will be registered as a module parameter
model = torch.jit.load('model.pt')结果是 32ms/image:
nn :forward: 32.35 (ms) 30.91 (imgs/s)
nn :backward: 0.00 (ms)
nn :update: 0.00 (ms)
nn :total: 32.35 (ms) 30.91 (imgs/s)以上过程对于 libtorch同样适用，也就是说通过python接口save下来的 .pt文件可以在C++中加载。
3. 环境变量设置
如果是single instance, 需要限制OpenMP thread数量以及CPU binding的方式，如下：
export OMP_NUM_THREADS=[number_of_physical_cores]
export KMP_AFFINITY=granularity=fine,compact,1,0如果是dual socket的CPU而只想用single socket跑，为了避免remote memory access，需要限制 numactrl：
numactl --physcpubind=0-$LAST_CORE --membind=0如果是multi instance，每个instance都会展开独立的OpenMP thread pool，也就是每个instance都需要限制 OMP_NUM_THREADS。保证 omp_threads * num_instances 不会超出物理核的数量，不然就会over subscription，就是抢核，这种情况对于Xeon来说是灾难性的。
multi instance 的情况要比single instance复杂很多，因为上层的threading model可能有很多种，可以是 torch.multiprocessing, std::threads, TBB等等。另外，设置affinity的方式要具体问题具体分析。社区上遇到很多这种multi instance没设置好环境变量的，问题五花八门，一般来讲分析此类问题最方便的就是用vtune，先查一查多少个omp master在同时跑。
4. 使用Intel OpenMP
目前PyTorch默认是用GNU OpenMP来编译的，就是libgomp.so。实际上Intel OpenMP也就是iomp5.so的性能要优于GNU版本。可以通过pre-loading的方式把gomp动态替换城iomp:
LD_PRELOAD=/opt/intel/compilers_and_libraries/linux/lib/intel64/libiomp5.so ./your_script.sh5. 使用jemalloc
PyTorch使用的是动态图，本身有一个缺点，就是output每次都是重新分配内存的。batch size比较大的时候memory allocation的开销是很大的（clear page），batch size为1的时候开销很小。 这个问题通过 jemalloc 可以得到一定程度的缓解，jemalloc是替换libc中的malloc的。
LD_PRELOAD=/home/mingfeim/packages/jemalloc-5.2.0/lib/libjemalloc.so ./your_script.sh个人经验jemalloc不一定每次都有效果，有个时候用vtune明明能看到很大的malloc开销，不过jemalloc cache不住，具体原因要等以后研究一下这里面的算法才知道。对于rnn batch_size = 1的情况非常有效，能提速差不多1/3。不管怎么样试一试，反正很简单，不花时间。
6. 用ICC编译PyTorch
PyTorch默认是用GCC编译的，不过GCC的auto vectorization非常差，不管是4.8.5还是7.1.0都完全无法和ICC相比，使用ICC对于没有手动展开的code会有很大提成。PyTorch 0.4版本之前是可以用ICC直接编译的，所以那时候intel/pytorch默认是用ICC的。不过和Caffe2的代码merge到一起之后用ICC编译有两千来个error和warning，可以把caffe2禁掉 BUILD_CAFFE2_OPS=0 。
CC=icc CXX=icpc python setup.py build7. DataLoader
torch.utils.data.DataLoader里面现在有个bug，我测试结果是num_workers > 0反而会变慢。实测的时候可以和num_workers = 0比较一下。TODO: 以后有时间处理掉。
8. Profile PyTorch
现在torch.autograd.profiler的功能已经做得很完善了，很容易就能找到workload中的hotspot。profiler的用法写在pytorch_profiler_parser。PyTorch里面还有一些operator在CPU是sequential的，发现问题可以提issue。
最近和外面的人接触比较多，很多人反馈PyTorch做deployment不方便，还有CPU上性能不好等等。其实PyTorch一直在往这个方向努力，也有相关解决方案，不过多少有点绕，所以写一下。用起来确实有点麻烦，不喜勿喷，都有难处。。。
