PyTorch code变动趋势是把TH开头这些模块逐渐往ATen native里面挪,native大概意思是pytorch重新写的部分,TH这些从lua torch继承来的称为legacy。大概从v0.3之后就是这个趋势,已经很长时间了。还有一个趋势就是python的code往c++中挪,比如cpu上面rnn的逻辑最开始都是.py的,现在都进c++了。 如果关注performance optimization,或者说需要改C++ code的话现在应该是关注ATen的,当然大多数人不会有这种需求。
最早移过去的是cudnn的部分,GPU上面Conv,BN和RNN都主要依赖这个模块。实现的时候有些trick,比如rnn的flatten weight等。TH,THC和THNN,THCUNN内容太多,一般来讲是如果需要优化那个operator就在aten里面重写直接bypass原有THXX中的部分。另外现在THS和THSC这两个sparse的包已经完全没有了,近期改掉的。
比如对于CPU的优化部分,原来TH的做法是尽量用替换TH_TENSOR_APPLY,这个宏是串行的,这个宏简直就是pytorch的原罪。向量化依赖于非常有限的THVector的逻辑,而且fb表现出对avx512非常抗拒,最多只有avx2的逻辑。现在aten中的做法完全不一样了,方便很多。
实际上如果只考虑功能性的话,用aten创建新的operator非常简单,如果新的operator可以用aten现有的operator实现的话,是不需要写反向的,autograd会自动生成反向代码,参考conv tbc的实现。如果前向操作不是用aten原有操作搭出来的,那需要提供相应的反向操作,然后告诉autograd自动求导的规则。这个设计是非常精髓的。
大多数人不需要动Module.cpp或者任何csrc下面的东西,除非你要加新的python接口,而且这个接口都不能被aten cover住,一般是一些global context的设置接口。计算的operator都是通过aten在c++实现然后bind到python上。
一般情况下也不会去动autograd的内容。如果需要添加新的operator,pytorch的做法是定义自动求导的规则,在derivatives.yaml里面,不需要知道autograd的实现细节。不过autograd目前有个问题是cpu上面的threading model, forward是和backward不是同一个process,导致结果就是会有两个omp thread pool,这个对peformance并不是十分友好,原则上要处理掉的。