笔者最近在尝试实现AlexNet的底层算子,基于pytorch的框架,本文主要记录一下pytorch中是如何实现relu算子的。
首先最外层是位于torch\nn\modules\activation.py
,主要代码如下:
py
__constants__ = ["inplace"]
inplace: bool
def __init__(self, inplace: bool = False):
super().__init__()
self.inplace = inplace
def forward(self, input: Tensor) -> Tensor:
return F.relu(input, inplace=self.inplace)
def extra_repr(self) -> str:
inplace_str = "inplace=True" if self.inplace else ""
return inplace_str
调用的是位于torch\nn\functional.py
的如下代码:
py
def relu(input: Tensor, inplace: bool = False) -> Tensor: # noqa: D400,D402
r"""relu(input, inplace=False) -> Tensor
Applies the rectified linear unit function element-wise. See
:class:`~torch.nn.ReLU` for more details.
"""
if has_torch_function_unary(input):
return handle_torch_function(relu, (input,), input, inplace=inplace)
if inplace:
result = torch.relu_(input)
else:
result = torch.relu(input)
return result
然后调用的是aten\src\ATen\native\Activation.cpp
的如下代码:
Tensor relu(const Tensor & self) {
TORCH_CHECK(self.scalar_type() != at::kBool, "Boolean inputs not supported for relu");
return at::clamp_min(self, 0);
}
可以看到,主要就是一个大小的比较。
pytorch调试工具
先说问题,只能看到python的处理逻辑,不能看到底层的C++的处理逻辑。
如何使用,参考的是这篇文章。注意,pdb虽然是python内置的包,但是仍然需要通过import pdb导入才能使用。
还有一个问题就是,pytorch是如何通过python代码调用C++代码的,留到下一篇博文更新。