分类：CUDA扩展

2020
09-25

PyTorch中的C++扩展实现

今天要聊聊用PyTorch进行C++扩展。在正式开始前，我们需要了解PyTorch如何自定义module。这其中，最常见的就是在python中继承torch.nn.Module，用PyTorch中已有的operator来组装成自己的模块。这种方式实现简单，但是，计算效率却未必最佳，另外，如果我们想实现的功能过于复杂，可能PyTorch中那些已有的函数也没法满足我们的要求。这时，用C、C++、CUDA来扩展PyTorch的模块就是最佳的选择了。由于目前市面上大... 继续阅读 >