2019-12 PyTorch: An Imperative Style, High-Performance Deep Learning Library
设计迎合4大趋势:
array-based (Tensor)
GPU加速
自动求导 (Auto Differentiation)
拥抱Python生态
4大设计原则:
使用算法和数据开发者熟悉的Python做编程接口,接入Python生态;
面向researcher;接口尽量简单,把复杂性隐藏起来;
性能别做的太差;为了易用性,牺牲10%性能是可以的,牺牲太多不行;
简单而不完善的设计:目的是把有限的人力,投入到跟上AI发展的快节奏,快速支持新feature上;
一切皆Python程序
易于扩展:
新模型更容易原生支持:(GAN网络例子)
方便在中间随时print中间变量(Tensor)的值,方便设断点debug;
兼容Python生态,易于扩展
torch.from_numpy()和Tensor.numpy(),其返回值和原值,是share同一片内存的,改一个的值另一个的值也会改动;
新增Module,只需要继承自Module类,写forward();
新增operator,如果其操作由PyTorch的operator组成,则直接写forward即可;如果包含PyTorch不支持的操作,需要使用torch.compile并写forward和backward;Python Custom Operators --- PyTorch Tutorials 2.3.0+cu121 documentation
新增Dataset,只需实现__getitem__和__len__; 其余的事交给DataLoader(shuffle, batch, 多进程, pin memory)
自动求导(AD)
采用reverse-mode automatic differentiation; <<Automatic differentiation in pytorch>>
forward之后,backward之前,如果中间环节的Tensor值被人为改写了(这些值在backward计算中会被用到),怎么办?答:Tensor使用版本号;
性能
Python的GIL锁:任何时刻,都只能有1个线程在执行;
核心代码使用C++,C++多线程并行执行,绕开了Python的GIL问题;
控制流和数据流,严格分开;控制流:if-else、循环等;数据流:op的实际执行;op在CPU侧只是launch kernel,kernel入队到CUDA stream里顺序执行,CPU侧和GPU侧是异步的;
痛点:cudaFree是同步操作,CPU侧会block住,等待所有stream里的kernels执行完,才释放显存;解决方案:自己维护显存池,reserved memory,少调用cudaFree;
one-pool-per-stream,没看懂;
Python的multiprocessing用的是磁盘来做进程通信,低效;PyTorch使用了torch.multiprocessing,用的是shared memory做进程通信;
Tensor的内存释放,使用引用计数,python侧和C++侧,都统计;
异步&overlap:
首轮迭代,cudaMalloc和cudaFree导致GPU使用率低下;第二轮开始,显存池发挥作用,cudaMalloc和cudaFree没有了,GPU使用率变高;