0 背景
给出一段代码,可以看到CNN是一个类,model = CNN()
为类的实例化。
问题:outputs = model(imgs)
一个类的实例,为什么能作为一个函数直接调用?并且调用的是forward方法?
...
class CNN(nn.Module):
def __init__(self):
super(CNN, self).__init__()
self.net = nn.Sequential(
# 卷积层
...
)
def forward(self, x):
return self.net(x)
...
model = CNN() # 模型实例化
...
for epoch in range(epochs):
train_loss = 0
for imgs, labels in train_dataloader:
imgs = imgs
labels = labels
outputs = model(imgs) # 前向计算
loss = loss_func(outputs, labels) # 损失函数计算
train_loss += loss.item() * imgs.size(0)
optimizer.zero_grad()
loss.backward()
optimizer.step() # 更新优化器
1 原码解析
版本信息:
python 3.10.0
pytorch 2.1.0
- torch.nn.Module的代码,在
torch\nn\modules\module.py
,下面这段原码是精简版的。
调用原理: python类的__call__方法支持实例可以作为函数来调用,即调用实例就是调用__call__
。
核心逻辑:nn.Module
的__call__
被定义为__wrapped_call_impl
,进而会调用到self.forward
。
...
class Module:
...
def __init__(self, *args, **kwargs) -> None:
...
__call__ : Callable[..., Any] = _wrapped_call_impl
...
def _wrapped_call_impl(self, *args, **kwargs):
if self._compiled_call_impl is not None:
return self._compiled_call_impl(*args, **kwargs) # type: ignore[misc]
else:
return self._call_impl(*args, **kwargs)
def _call_impl(self, *args, **kwargs):
forward_call = (self._slow_forward if torch._C._get_tracing_state() else self.forward)
# If we don't have any hooks, we want to skip the rest of the logic in
# this function, and just call forward.
if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
or _global_backward_pre_hooks or _global_backward_hooks
or _global_forward_hooks or _global_forward_pre_hooks):
return forward_call(*args, **kwargs)
...
2 引申
有了这个核心逻辑,我们也可以实现一个简单的调用。
这里,test
方法可以类比于nn.Module
的forward
方法。
另外,在Python中,
Callable[..., Any]
是一个类型注解,用来描述一个可以被调用的对象。
这里的语法解释如下:
Callable
是一个类型,表示一个可调用的类型。
[...]
表示参数类型是可变的。这意味着可以传递任意数量的参数给这个可调用的对象。
Any
表示返回类型是任意的,即这个可调用的对象可以返回任何类型的结果。
from typing import Any, Callable
class Module:
def __init__(self):
return
def test(self, *args):
print("success:", sum(args))
return
__call__ : Callable[..., Any] = test
A = Module()
output = A(5,6,7)
输出
success: 18