pytorch 中_call_impl()函数

记录pytorch 版本中的 nn.Module() 重要函数

1. _call_impl()

1.1 torch1.7.1 版本

python 复制代码
    def _call_impl(self, *input, **kwargs):
        for hook in itertools.chain(
                _global_forward_pre_hooks.values(),
                self._forward_pre_hooks.values()):
            result = hook(self, input)
            if result is not None:
                if not isinstance(result, tuple):
                    result = (result,)
                input = result
        if torch._C._get_tracing_state():
            result = self._slow_forward(*input, **kwargs)
        else:
            result = self.forward(*input, **kwargs)
        for hook in itertools.chain(
                _global_forward_hooks.values(),
                self._forward_hooks.values()):
            hook_result = hook(self, input, result)
            if hook_result is not None:
                result = hook_result
        if (len(self._backward_hooks) > 0) or (len(_global_backward_hooks) > 0):
            var = result
            while not isinstance(var, torch.Tensor):
                if isinstance(var, dict):
                    var = next((v for v in var.values() if isinstance(v, torch.Tensor)))
                else:
                    var = var[0]
            grad_fn = var.grad_fn
            if grad_fn is not None:
                for hook in itertools.chain(
                        _global_backward_hooks.values(),
                        self._backward_hooks.values()):
                    wrapper = functools.partial(hook, self)
                    functools.update_wrapper(wrapper, hook)
                    grad_fn.register_hook(wrapper)
        return result

以上的函数的功能作用解释如下:

提供的代码是 PyTorch 模块方法的 _call_impl 实现。当模块用作可调用对象时,通常使用输入数据调用它时,将调用此方法。让我们逐步分解代码以了解其功能:
for hook in itertools.chain(_global_forward_pre_hooks.values(), self._forward_pre_hooks.values()):

此循环遍历两个钩子集合:

_global_forward_pre_hooks 和 _forward_pre_hooks 。

钩子是可以注册为在神经网络向前或向后传递期间在特定点执行的函数。

这些预置挂钩将在模块的实际前向传递之前执行。

result = hook(self, input):

对于每个钩子,使用 self (模块)和 input 参数调用 hook 该函数。

if result is not None:

如果钩子返回非 None 值,则表示钩子修改了输入数据,并且此修改后的数据将成为循环中下一个挂钩的新输入。

if not isinstance(result, tuple): result = (result,)

钩子的结果将转换为元组(如果它还没有元组)。这是为了处理钩子可能返回单个值而不是元组的情况。

input = result

修改后的输入将成为循环中下一个钩子的新输入。
if torch._C._get_tracing_state(): result = self._slow_forward(*input, **kwargs)

此条件块检查是否正在跟踪当前正向传递。如果是,则使用修改后的输入数据调用该方_slow_forward 。

else: result = self.forward(*input, **kwargs)

如果未跟踪正向传递,则使用修改后的输入数据调用模块的正常 forward 方法。
for hook in itertools.chain(_global_forward_hooks.values(), self._forward_hooks.values()):

此循环遍历两个前向钩子集合: _global_forward_hooks 和 _forward_hooks 。这些钩子在模块的正向传递之后执行。

hook_result = hook(self, input, result):

对于每个钩子,使用 self 、 input 和 result 作为参数调用 hook 函数。

if hook_result is not None: result = hook_result

如果钩子返回非 None 值,则表示钩子修改了正向传递的结果,并且此修改后的结果将成为循环中下一个挂钩的新结果。
if (len(self._backward_hooks) > 0) or (len(_global_backward_hooks) > 0):

此条件块检查是否有任何全局或为此特定模块注册的向后钩子。

var = result

正向传递的结果存储在变量 var 中。
while not isinstance(var, torch.Tensor):

此循环迭代 till 是 var 的 torch.Tensor 实例。

if isinstance(var, dict): var = next((v for v in var.values() if isinstance(v, torch.Tensor)))

如果是一个字典,它 var 查找字典中的第一个值,即 torch.Tensor .

else: var = var[0]

如果不是字典,它 var 假定它是一个序列(例如,列表,元组)并获取其第一个元素。

grad_fn = var.grad_fn

grad_fn torch.Tensor 实例的属性被分配给变量 grad_fn 。此属性表示在反向传播期间计算张量梯度的函数。
if grad_fn is not None:

如果 不是 grad_fn None,则表示张量参与了需要梯度的计算,我们需要向其附加向后钩子。

for hook in itertools.chain(_global_backward_hooks.values(), self._backward_hooks.values()):

此循环遍历两个向后钩子集合: _global_backward_hooks 和 _backward_hooks 。这些钩子在模块的向后传递期间执行。

wrapper = functools.partial(hook, self)

对于每个向后钩子,通过将钩子函数与模块 self 作为参数部分应用来创建一个新函数。这样做是为了确保钩子函数可以访问模块。

functools.update_wrapper(wrapper, hook)

包装器函数使用原始钩子函数中的信息进行更新,例如其名称和文档字符串。

grad_fn.register_hook(wrapper)

包装器函数注册为挂接到 grad_fn .这意味着在反向传播期间计算梯度时,将执行钩子以对梯度执行其他操作。

return result

返回前向传递的最终结果。

总之,该方法 _call_impl 执行前向挂钩(前向前和后向前),执行前向传递,执行向后挂钩(如果需要),并返回向前传递的结果。它还处理钩子修改输入或结果数据的情况,并确保向后钩子附加到相关张量,以便在反向传播期间进行梯度计算。

相关推荐
说私域9 分钟前
数字经济时代下的创新探索与实践:以“开源AI智能名片2+1链动模式S2B2C商城小程序源码”为核心
人工智能·小程序·开源
游客52025 分钟前
设计模式-结构型-装饰器模式
python·设计模式·装饰器模式
petunsecn26 分钟前
将 AzureBlob 的日志通过 Azure Event Hubs 发给 Elasticsearch(3 纯python的经济方案)
python·elasticsearch·microsoft·azure
测试杂货铺41 分钟前
单元测试与unittest框架
自动化测试·软件测试·python·测试工具·职场和发展·单元测试·测试用例
伪_装1 小时前
[Bug]libGL.so.1: cannot open shared object file: No such file or directory
人工智能·opencv·bug
思码逸研发效能1 小时前
软件工程3.0时代,AI落地研效成熟时
人工智能·aigc·研发效能·ai编程·敏捷开发·研发效能度量
yzx9910131 小时前
Python根据图片生成学生excel成绩表
python·mysql·excel
深图智能1 小时前
PyTorch使用教程(4)-torch.nn
人工智能·pytorch·深度学习
Channing Lewis1 小时前
Python 3.9及以上版本支持的新的字符串函数 str.removeprefix()
服务器·python
一语雨在生无可恋敲代码~1 小时前
wandb使用遇到的一些问题
pytorch·计算机视觉