跳至主要內容

Torch

Genhiy...大约 15 分钟PythonCodeBook

tensor

zeros()、ones()

创建全0、全1的tensor:

>>> torch.zeros([2, 4], dtype=torch.int32)
tensor([[ 0,  0,  0,  0],
        [ 0,  0,  0,  0]], dtype=torch.int32)
>>> cuda0 = torch.device('cuda:0')
>>> torch.ones([2, 4], dtype=torch.float64, device=cuda0)
tensor([[ 1.0000,  1.0000,  1.0000,  1.0000],
        [ 1.0000,  1.0000,  1.0000,  1.0000]], dtype=torch.float64, device='cuda:0')

item()

从只有一个数字的tensor中得到这个数字:

>>> x = torch.tensor([[1]])
>>> x
tensor([[ 1]])
>>> x.item()
1
>>> x = torch.tensor(2.5)
>>> x
tensor(2.5000)
>>> x.item()
2.5

grad、backward()

获取梯度与反向传播:

>>> x = torch.tensor([[1., -1.], [1., 1.]], requires_grad=True)
>>> out = x.pow(2).sum()
>>> out.backward()
>>> x.grad
tensor([[ 2.0000, -2.0000],
        [ 2.0000,  2.0000]])

unsqueeze()

升维,使用方法:input.unsqueeze(1),1表示在张量的第二维中加括号,也即在一个torch.Size([8, 6, 224, 224])的第二维上插一个维度变为:torch.Size([8, 1, 6, 224, 224])

repeat()

重复,关于这个的操作,有些文章写的那叫一个离谱,根本就读不懂,其实很简单,首先,repeat中的参数个数不能少于张量的维度,当等于时,张量的第i个维度重复repeat的第i个参数倍次,当repeat参数个数大于张量维度时,就在张量维度上补1,补至和repeat的参数个数相同,然后再进行repeat操作。

input = torch.rand(8,6,224,224)
input = input.repeat(2,2,1,1)
input.shape
>>> torch.Size([16, 12, 224, 224])
input = input.repeat(1,1,2,1,1)
input.shape
>>> torch.Size([2, 8, 12, 224, 224])

torch.flip()

用于翻转张量的函数。它可以用于在指定维度上对张量进行翻转操作。

torch.flip(input, dims) → Tensor
  • input:输入张量,可以是任意形状的张量。
  • dims:一个整数或整数列表,表示要翻转的维度。
  • 返回一个张量,表示在指定维度上翻转后的结果,dims=1代表将第一维度反转,dims=[1,2]代表将第一、二维度反转。

model

nn.Module

定义:

class Module(object):
    def __init__(self): # 定义一些基本内容,需重写,
    def forward(self, *input): # 前向过程,需重写
 
    def add_module(self, name, module):
    def cuda(self, device=None): # 将模型移到cuda设备上
    def cpu(self): # 将模型移到cpu上
    def __call__(self, *input, **kwargs): # model(x)这条语句可以正常执行的前提,这个函数调用的forward
    def parameters(self, recurse=True): # 查看神经网络的参数信息
    def named_parameters(self, prefix='', recurse=True): # 同样是查看神经网络的参数信息
    def children(self): # 返回网络模型里的组成元素,children()返回的是最外层的元素
    def named_children(self):
    def modules(self): # 返回网络模型里的组成元素,modules()返回的是所有的元素,包括不同级别的子元素
    def named_modules(self, memo=None, prefix=''):
    def train(self, mode=True): # 设置模式为训练
    def eval(self): # 设置模式为评估
    def zero_grad(self):
    def __repr__(self):
    def __dir__(self):
'''
有一部分没有完全列出来
'''

一般把网络中具有可学习参数的层(如全连接层、卷积层等)放在构造函数__init__()中,当然也可以把不具有参数的层也放在里面;如果这些不含参数的层(如ReLU、dropout、BatchNormanation层)不放在构造函数__init__里面,则在forward方法里面可以使用nn.functional来代替:

import torch
import torch.nn.functional as F
 
class MyNet(torch.nn.Module):
    def __init__(self):
        super(MyNet, self).__init__()  # 第一句话,调用父类的构造函数
        self.conv1 = torch.nn.Conv2d(3, 32, 3, 1, 1)
        self.conv2 = torch.nn.Conv2d(3, 32, 3, 1, 1)
        self.dense1 = torch.nn.Linear(32 * 3 * 3, 128)
        self.dense2 = torch.nn.Linear(128, 10)
 
    def forward(self, x):
        x = self.conv1(x)
        x = F.relu(x)
        x = F.max_pool2d(x)
        x = self.conv2(x)
        x = F.relu(x)
        x = F.max_pool2d(x)
        x = self.dense1(x)
        x = self.dense2(x)
        return x
 
model = MyNet()
print(model)
'''运行结果为:
MyNet(
  (conv1): Conv2d(3, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (conv2): Conv2d(3, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (dense1): Linear(in_features=288, out_features=128, bias=True)
  (dense2): Linear(in_features=128, out_features=10, bias=True)
)
'''

注意:此时,将没有训练参数的层没有放在构造函数里面了,所以这些层就不会出现在model里面,但是运行关系是在forward里面通过functional的方法实现的。

其他构建网络的方法:

import torch.nn as nn
from collections import OrderedDict
class MyNet(nn.Module):
    def __init__(self):
        super(MyNet, self).__init__()
        self.conv_block = nn.Sequential(
            nn.Conv2d(3, 32, 3, 1, 1),
            nn.ReLU(),
            nn.MaxPool2d(2))
        self.dense_block = nn.Sequential(
            nn.Linear(32 * 3 * 3, 128),
            nn.ReLU(),
            nn.Linear(128, 10)
        )
    # 在这里实现层之间的连接关系,其实就是所谓的前向传播
    def forward(self, x):
        conv_out = self.conv_block(x)
        res = conv_out.view(conv_out.size(0), -1)
        out = self.dense_block(res)
        return out
 
model = MyNet()
print(model)
'''运行结果为:
MyNet(
  (conv_block): Sequential(
    (0): Conv2d(3, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): ReLU()
    (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  )
  (dense_block): Sequential(
    (0): Linear(in_features=288, out_features=128, bias=True)
    (1): ReLU()
    (2): Linear(in_features=128, out_features=10, bias=True)
  )
)
'''

这里在每一个包装块里面,各个层是没有名称的,默认按照0、1、2、3、4来排名。如果我们希望有名字:

import torch.nn as nn
from collections import OrderedDict
class MyNet(nn.Module):
    def __init__(self):
        super(MyNet, self).__init__()
        self.conv_block = nn.Sequential(
            OrderedDict(
                [
                    ("conv1", nn.Conv2d(3, 32, 3, 1, 1)),
                    ("relu1", nn.ReLU()),
                    ("pool", nn.MaxPool2d(2))
                ]
            ))
 
        self.dense_block = nn.Sequential(
            OrderedDict([
                ("dense1", nn.Linear(32 * 3 * 3, 128)),
                ("relu2", nn.ReLU()),
                ("dense2", nn.Linear(128, 10))
            ])
        )
 
    def forward(self, x):
        conv_out = self.conv_block(x)
        res = conv_out.view(conv_out.size(0), -1)
        out = self.dense_block(res)
        return out
 
model = MyNet()
print(model)
'''运行结果为:
MyNet(
  (conv_block): Sequential(
    (conv1): Conv2d(3, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (relu1): ReLU()
    (pool): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  )
  (dense_block): Sequential(
    (dense1): Linear(in_features=288, out_features=128, bias=True)
    (relu2): ReLU()
    (dense2): Linear(in_features=128, out_features=10, bias=True)
  )
)
'''

以及简洁一些:

import torch.nn as nn
from collections import OrderedDict
class MyNet(nn.Module):
    def __init__(self):
        super(MyNet, self).__init__()
        self.conv_block=torch.nn.Sequential()
        self.conv_block.add_module("conv1",torch.nn.Conv2d(3, 32, 3, 1, 1))
        self.conv_block.add_module("relu1",torch.nn.ReLU())
        self.conv_block.add_module("pool1",torch.nn.MaxPool2d(2))
 
        self.dense_block = torch.nn.Sequential()
        self.dense_block.add_module("dense1",torch.nn.Linear(32 * 3 * 3, 128))
        self.dense_block.add_module("relu2",torch.nn.ReLU())
        self.dense_block.add_module("dense2",torch.nn.Linear(128, 10))
 
    def forward(self, x):
        conv_out = self.conv_block(x)
        res = conv_out.view(conv_out.size(0), -1)
        out = self.dense_block(res)
        return out
 
model = MyNet()
print(model)
'''运行结果为:
MyNet(
  (conv_block): Sequential(
    (conv1): Conv2d(3, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (relu1): ReLU()
    (pool1): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  )
  (dense_block): Sequential(
    (dense1): Linear(in_features=288, out_features=128, bias=True)
    (relu2): ReLU()
    (dense2): Linear(in_features=128, out_features=10, bias=True)
  )
)
'''

保存和加载模型

# Define model
class TheModelClass(nn.Module):
    def __init__(self):
        super(TheModelClass, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

# Initialize model
model = TheModelClass()

# Initialize optimizer
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

# Print model's state_dict
print("Model's state_dict:")
for param_tensor in model.state_dict():
    print(param_tensor, "\t", model.state_dict()[param_tensor].size())

# Print optimizer's state_dict
print("Optimizer's state_dict:")
for var_name in optimizer.state_dict():
    print(var_name, "\t", optimizer.state_dict()[var_name])

上述代码先是简单定义一个 5 层的 CNN,然后分别打印模型的参数和优化器参数。

输出结果:

Model's state_dict:
conv1.weight     torch.Size([6, 3, 5, 5])
conv1.bias   torch.Size([6])
conv2.weight     torch.Size([16, 6, 5, 5])
conv2.bias   torch.Size([16])
fc1.weight   torch.Size([120, 400])
fc1.bias     torch.Size([120])
fc2.weight   torch.Size([84, 120])
fc2.bias     torch.Size([84])
fc3.weight   torch.Size([10, 84])
fc3.bias     torch.Size([10])

Optimizer's state_dict:
state    {}
param_groups     [{'lr': 0.001, 'momentum': 0.9, 'dampening': 0, 'weight_decay': 0, 'nesterov': False, 'params': [4675713712, 4675713784, 4675714000, 4675714072, 4675714216, 4675714288, 4675714432, 4675714504, 4675714648, 4675714720]}]

查看model各层输入输出

提示

本内容将极大协助您接触一个新的模型。

import timm
import torch
from torchinfo import summary

model = timm.create_model(
    'efficientvit_b0.r224_in1k',
    pretrained=True,
    features_only=True,
    pretrained_cfg_overlay=dict(file='/mnt/iag/user/daiyiheng/ckpt/backbones/efficientvit_b0/pytorch_model.bin'),
)
model = model.eval()
model.cuda()

img_size = (4, 3, 224, 480)
summary(model, img_size)

x = torch.randn((4, 3, 224, 480))
x = x.cuda()
for o in model(x):
    print(o.shape)

输出结果:

===================================================================================================
Layer (type:depth-idx)                                  Output Shape              Param #
===================================================================================================
FeatureListNet                                    [4, 16, 56, 120]          --
├─ConvNormAct: 1-1                                [4, 8, 112, 240]          --
│    └─Dropout: 2-1                               [4, 3, 224, 480]          --
│    └─Conv2d: 2-2                                [4, 8, 112, 240]          216
│    └─BatchNorm2d: 2-3                           [4, 8, 112, 240]          16
│    └─Hardswish: 2-4                             [4, 8, 112, 240]          --
├─ResidualBlock: 1-2                              [4, 8, 112, 240]          --
│    └─Identity: 2-5                              [4, 8, 112, 240]          --
│    └─DSConv: 2-6                                [4, 8, 112, 240]          --
│    │    └─ConvNormAct: 3-1                      [4, 8, 112, 240]          88
│    │    └─ConvNormAct: 3-2                      [4, 8, 112, 240]          80
│    └─Identity: 2-7                              [4, 8, 112, 240]          --
├─EfficientVitStage: 1-3                          [4, 16, 56, 120]          --
│    └─Sequential: 2-8                            [4, 16, 56, 120]          --
│    │    └─ResidualBlock: 3-3                    [4, 16, 56, 120]          1,216
│    │    └─ResidualBlock: 3-4                    [4, 16, 56, 120]          2,912
├─EfficientVitStage: 1-4                          [4, 32, 28, 60]           --
│    └─Sequential: 2-9                            [4, 32, 28, 60]           --
│    │    └─ResidualBlock: 3-5                    [4, 32, 28, 60]           3,968
│    │    └─ResidualBlock: 3-6                    [4, 32, 28, 60]           9,920
├─EfficientVitStage: 1-5                          [4, 64, 14, 30]           --
│    └─Sequential: 2-10                           [4, 64, 14, 30]           --
│    │    └─ResidualBlock: 3-7                    [4, 64, 14, 30]           13,824
│    │    └─EfficientVitBlock: 3-8                [4, 64, 14, 30]           64,192
│    │    └─EfficientVitBlock: 3-9                [4, 64, 14, 30]           64,192
├─EfficientVitStage: 1-6                          [4, 128, 7, 15]           --
│    └─Sequential: 2-11                           [4, 128, 7, 15]           --
│    │    └─ResidualBlock: 3-10                   [4, 128, 7, 15]           52,224
│    │    └─EfficientVitBlock: 3-11               [4, 128, 7, 15]           234,880
│    │    └─EfficientVitBlock: 3-12               [4, 128, 7, 15]           234,880
===================================================================================================
Total params: 682,608
Trainable params: 682,608
Non-trainable params: 0
Total mult-adds (Units.MEGABYTES): 764.16
===================================================================================================
Input size (MB): 5.16
Forward/backward pass size (MB): 317.40
Params size (MB): 2.73
Estimated Total Size (MB): 325.29
=========================================================================================================
torch.Size([4, 16, 56, 120])
torch.Size([4, 32, 28, 60])
torch.Size([4, 64, 14, 30])
torch.Size([4, 128, 7, 15])

我们可以很直观的从返回的信息中看到:本网络是一个FeatureListNet,其中包括第一层级模块6个,分别是一个ConvNormAct、一个ResidualBlock和四个EfficientVitStage,但是四个EfficientVitStage略微有些区别。

而且我们还可以直观的看到各个层的参数数目、总参数数目以及可训练的参数数目。以及总站用的内存数量:325.29MB。

这里efficientvit_b0会是有四个输出,如输出内容的最下方显示,分别是:torch.Size([4, 16, 56, 120])torch.Size([4, 32, 28, 60])torch.Size([4, 64, 14, 30])torch.Size([4, 128, 7, 15])。但是上方FeatureListNet的输出仅显示torch.Size([4, 16, 56, 120]),原因是多输出只显示最先输出的一个。

预训练模型读取

上面展示了用timm中读取预训练模型:

model = timm.create_model(
    'efficientvit_b0.r224_in1k',
    pretrained=True,
    features_only=True,
    pretrained_cfg_overlay=dict(file='/mnt/iag/user/daiyiheng/ckpt/backbones/efficientvit_b0/pytorch_model.bin'),
)

下面是用torchvision读取:

model = torchvision.models.efficientnet_b4()
model.load_state_dict(torch.load('/mnt/iag/user/daiyiheng/ckpt/backbones/efficientnet_b4_rwightman-23ab8bcd.pth'))

nn操作层

上采样层nn.Upsample

CLASS torch.nn.Upsample(size=None, scale_factor=None, mode='nearest', align_corners=None)

具体参数及用法:pytorch torch.nn 实现上采样——nn.Upsampleopen in new window

上采样一个给定的多通道的 1D (temporal,如向量数据), 2D (spatial,如jpg、png等图像数据) or 3D (volumetric,如点云数据)数据。

假设输入数据的格式为minibatch x channels x [optional depth] x [optional height] x width。因此对于一个空间spatial输入,我们期待着4D张量的输入,即minibatch x channels x height x width。而对于体积volumetric输入,我们则期待着5D张量的输入,即minibatch x channels x depth x height x width。

对于上采样有效的算法分别有对 3D, 4D和 5D 张量输入起作用的 最近邻、线性,、双线性, 双三次(bicubic)和三线性(trilinear)插值算法。

恒等映射nn.Identity()

nn.Identity()实际上是一个恒等映射,不对输入进行任何变换或操作,只是简单地将输入返回作为输出。

既然nn.Identity() 是 PyTorch 中的一个层(layer)。它实际上是一个恒等映射,不对输入进行任何变换或操作,只是简单地将输入返回作为输出。那为什么要写它呢?它存在的意义是什么?nn.Identity() 层在神经网络中可能看起来似乎没有实际用处,因为它不对输入进行任何操作或变换,只是简单地将输入返回作为输出。然而,它实际上在某些情况下是很有用的,其存在意义包括以下几点:

  • 连接路径或跳跃连接(Skip Connections): 在一些网络架构中,特别是残差网络(Residual Networks)或者一些需要跨层连接的架构(比如 DenseNet),nn.Identity() 可以用于实现跳跃连接。通过跳跃连接,神经网络可以直接从一个层传递信息到后续层,这有助于解决梯度消失或梯度爆炸等问题,同时也有助于提高模型的性能和训练速度。
  • 模型组合和特殊结构设计: 在设计复杂的神经网络结构时,nn.Identity() 可以用于在模型中创建一些特殊的连接或分支结构。通过使用恒等映射,可以更轻松地实现某些复杂模型的组合,或者通过条件语句动态地选择是否应用某些层。
  • 代码一致性和灵活性: 在编写神经网络代码时,有时需要保持一致性,可能会需要一个占位符层来代表某些特定的操作。nn.Identity() 可以填补这个需求,即使不对输入进行任何更改,也能保持代码的一致性和清晰度。
  • 简化模型和调试: 在一些情况下,为了简化模型或者调试网络结构,可以使用 nn.Identity() 层。它允许将某些部分固定为恒等映射,方便单独地测试网络的不同部分。

functional方法

数组归一化 F.normalize

torch.nn.functional.normalize(input, p=2.0, dim=1, eps=1e-12, out=None)

功能:利用LpL_p范数对输入的数组沿特定的维度进行归一化:

对于尺寸为(n0,,ndim,,nk)(n_0,…,n_{dim},…,n_k)的输入数组input,每个ndimn_{dim}上的元素向量vv沿着维度dim进行转换,转换公式为:

v=vmax(vp,ϵ) v=\frac{v}{\max(||v||_p,\epsilon)}

范数计算公式

对于数据 x=[x1,x2,,xn]Tx=[x_1,x_2,\dots,x_n]^T:

Lp范数:xp=(x1p+x2p++xnp)1pL1范数:x1=x1+x2++xnL2范数:x2=(x12+x22++xn2)12 \begin{aligned}&\bullet L_{p} \text{范数:} ||x||_{p} = (|x_{1} |^{p}+|x_{2} |^{p}+\cdots+|x_{n} |^{p})^{\frac{1}{p}}\\&\bullet L_{1} \text{范数:} ||x||_{1} = |x_{1} |+|x_{2} |+\cdots+|x_{n} |\\&\bullet L_{2} \text{范数:} ||x||_{2} = (|x_{1} |^{2}+|x_{2} |^{2}+\cdots+|x_{n} |^{2})^{\frac{1}{2}}\end{aligned}

输入:

  • input:输入的数组,数组数据类型为float
  • p:指定使用的范数,数据类型为float,默认2.0
  • dim:指定的维度,数据类型为int,默认1
  • eps:边界值,防止分母为0,默认1e-12

本内容参考自:PyTorch学习笔记:F.normalize——数组归一化运算open in new window

debug相关方法

hooks方法

为了节省显存(内存),pytorch在计算过程中不保存中间变量,包括中间层的特征图和非叶子张量的梯度等。有时对网络进行分析时需要查看或修改这些中间变量,此时就需要注册一个钩子(hook)来导出需要的中间变量。

hook方法有四种:

  • torch.Tensor.register_hook()
  • torch.nn.Module.register_forward_hook()
  • torch.nn.Module.register_backward_hook()
  • torch.nn.Module.register_forward_pre_hook()

torch.Tensor.register_hook(hook)

用来导出指定张量的梯度,或修改这个梯度值。

import torch
def grad_hook(grad):
    grad *= 2
x = torch.tensor([2., 2., 2., 2.], requires_grad=True)
y = torch.pow(x, 2)
z = torch.mean(y)
h = x.register_hook(grad_hook)
z.backward()
print(x.grad)
h.remove()    # removes the hook
>>> tensor([2., 2., 2., 2.])

注意:

  • 上述代码是有效的,但如果写成 grad = grad * 2就失效了,因为此时没有对grad进行本地操作,新的grad值没有传递给指定的梯度。保险起见,最好在def语句中写明return grad。即:

    def grad_hook(grad):
        grad = grad * 2
        return grad
    
  • 可以用remove()方法取消hook。注意remove()必须在backward()之后,因为只有在执行backward()语句时,pytorch才开始计算梯度,而在x.register_hook(grad_hook)时它仅仅是"注册"了一个grad的钩子,此时并没有计算,而执行remove就取消了这个钩子,然后再backward()时钩子就不起作用了。

torch.nn.Module.register_forward_hook(module, inp, out)

用来导出指定子模块(可以是层、模块等nn.Module类型)的输入输出张量,但只可修改输出,常用来导出或修改卷积特征图。

简单范例:

inps, outs = [],[]
def layer_hook(module, inp, out):
    inps.append(inp[0].data.cpu().numpy())
    outs.append(out.data.cpu().numpy())

hook = net.layer1.register_forward_hook(layer_hook)
output = net(input)
hook.remove()

再具体些,比如有个LeNet:

import torch
import torch.nn as nn
import torch.nn.functional as F
 
class LeNet(nn.Module):
    def __init__(self):
        super(LeNet, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16*5*5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)
 
    def forward(self, x):
        out = self.conv1(x)
        out = F.relu(out)     
        out = F.max_pool2d(out, 2)      
        
        out = self.conv2(out)
        out = F.relu(out)  
        out = F.max_pool2d(out, 2)
        
        out = out.view(out.size(0), -1)
        out = F.relu(self.fc1(out))
        out = F.relu(self.fc2(out))
        out = self.fc3(out)
        return out

如果我们要获取conv2的输出,一种最直观的思路是直接修改forward部分的代码,将conv2的中间结果return即可:

def forward(self, x):
    out = self.conv1(x)
    out = F.relu(out)     
    out = F.max_pool2d(out, 2)      
    
    out = self.conv2(out)
    out_conv2 = out
    out = F.relu(out)
    out = F.max_pool2d(out, 2)
    
    out = out.view(out.size(0), -1)
    out = F.relu(self.fc1(out))
    out = F.relu(self.fc2(out))
    out = self.fc3(out)
    return out, out_conv2

但很多时候,我们并没有办法去直接修改网络的源代码,比如在pytorch中已经封装好的网络,那么这个时候就可以利用hook从外部获取Module的中间输出结果了。即:

features = []
def hook(module, input, output): 
    print(module,input[0].shape,output[0].shape)
    features.append(output.clone().detach())
 
net = LeNet() 
x = torch.randn(2, 3, 32, 32)  
handle = net.conv2.register_forward_hook(hook)
y = net(x)
print(features[0])
handle.remove()

输出结果:

Conv2d(6, 16, kernel_size=(5, 5), stride=(1, 1)) 
torch.Size([2, 6, 14, 14]) torch.Size([16, 10, 10])

取出网络的相应层后,对该层调用register_forward_hook方法。这个方法需要传入一个hook方法:

hook(module, input, output) -> None or modified output
  • module:表示该层网络
  • input:该层网络的输入
  • output:该层网络的输出

从这里可以发现hook甚至可以更改输入输出(不过并不会影响网络forward的实际结果),不过在这里我们只是简单地将output给保存下来。

需要注意的是hook函数在使用后应及时删除,以避免每次都运行增加运行负载。

参考资料:【PyTorch】 register_forward_hook()简单用法open in new window

torch.nn.Module.register_forward_pre_hook(module, in)

用来导出或修改指定子模块的输入张量。

def pre_hook(module, inp):
    inp0 = inp[0]
    inp0 = inp0 * 2
    inp = tuple([inp0])
    return inp

hook = net.layer1.register_forward_pre_hook(pre_hook)
output = net(input)
hook.remove()

注意:

  • inp值是个tuple类型,所以需要先把其中的张量提取出来,再做其他操作,然后还要再转化为tuple返回。
  • 在执行output = net(input)时才会调用此句,remove()可放在调用后用来取消钩子。

torch.nn.Module.register_backward_hook(module, grad_in, grad_out)

用来导出指定子模块的输入输出张量的梯度,但只可修改输入张量的梯度(即只能返回gin),输出张量梯度不可修改。

gouts = []
def backward_hook(module, gin, gout):
    print(len(gin),len(gout))
    gouts.append(gout[0].data.cpu().numpy())
    gin0,gin1,gin2 = gin
    gin1 = gin1*2
    gin2 = gin2*3
    gin = tuple([gin0,gin1,gin2])
    return gin

hook = net.layer1.register_backward_hook(backward_hook)
loss.backward()
hook.remove()

注意:

  • 其中的grad_in和grad_out都是tuple,必须要先解开,修改时执行操作后再重新放回tuple返回。
  • 这个钩子函数在backward()语句中被调用,所以remove()要放在backward()之后用来取消钩子。

参考资料:【pytorch学习】四种钩子方法(register_forward_hook等)的用法和注意点open in new window