Learning record 4:einops / / cudnn benchamark=true // hook

einops 

import torch
from einops import rearrange,reduce,repeat

x= torch.randn(2,3,8,8)
#1  Transpose operation 
out1 = x.transpose(1,2)
out2 = rearrange(x,'b c h w ->b h c w')
print('verify out1 & out2  ---->:',torch.allclose(out1,out2))

#2  deformation 
out3 = x.reshape(6,8,8)
out4 = rearrange(x,'b c h w -> (b c) h w')
x_restore = rearrange(out4,'(b c) h w -> b c h w ',b =2)
print('verify out3 & out4  ---->:',torch.allclose(out3,out4),'|| verify x & x_restore---->:',torch.allclose(x,x_restore))

#3 image2patch
out5= rearrange(x,'b c (h1 p1) (w1 p2) -> b c (h1 w1) (p1 p2)',p1=2,p2=2)  #  This got patch  yes non-overlapping Of 
print('out5 ---->:',x.size(),out5.size())
out6 = rearrange(out5,'b c n a -> b n (a c)')  # [batchsize, num_of_patches, patches_depth]
print('out6 ---->:',out6.size())

#4  Average pooling 
out7 = reduce(x,'b c h w -> b c','mean')
print('out7---->:',out7.size())

#5  The stack tensor
x_list = [x,x,x]
out8 = rearrange(x_list,'n b c h w -> n b c h w ')
print('out8---->:',out8.size())

#6  Dimension expansion 
out9 = rearrange(x,'b c h w -> b c h w 1 ') # Be similar to  torch.unsqueeze
#
print('out9---->:',out9.size())

#7  Copy 
out10 = repeat(out9,'b c h w 1 ->  b c h w 2  ')  # Be similar to  torch.tile
out11 = repeat(x,'b c h w -> b (2 c) h w ')  # Copy along the channel 
print('out10---->:',out10.size(),'out11---->:',out11.size())

The results are as follows :

torch.backends.cudnn.benchmark=True

Set up  torch.backends.cudnn.benchmark=True  It will take a little extra time for the program to start , Search the most suitable convolution algorithm for each convolution layer of the whole network , And then realize the acceleration of network . The applicable scenario is that the network structure is fixed （ It's not dynamic ）, Input shape of network （ Include batch size, Picture size , Input channel ） It is the same. , In fact, it is generally applicable . conversely , If the setting of the convolution layer keeps changing , Will cause the program to optimize continuously , It takes more time .

In fact, it is usually added at the beginning , For example, you can use GPU At the same time , I'll add a sentence later ：

if args.use_gpu and torch.cuda.is_available():
    device = torch.device('cuda')
    torch.backends.cudnn.benchmark = True
else:
    device = torch.device('cpu')

For when to set  torch.backends.cudnn.benchmark=True, A word is ： If the convolution network structure is not dynamic , Network input (batch size, The size of the image , Input channel ) Is constant , Then feel free to use .

hook Get the middle feature layer

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.conv3 = nn.Conv2d(16, 32, 5)

    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 = self.conv3(out)
        out = F.max_pool2d(out, 2)
        return out

# The first is to modify the network structure , Through the network return  Return the desired variable 

class LeNet_multi_outputs(nn.Module):
    def __init__(self):
        super(LeNet_multi_outputs, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.conv3 = nn.Conv2d(16, 32, 5)

    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 = self.conv3(out)
        out = F.max_pool2d(out, 2)
        return out,out_conv2


if __name__ == "__main__":
    model = LeNet_multi_outputs()
    print(model)
    print(model.conv1)
    input = torch.randn(1, 3, 224, 224)
    output,out_conv2 = model(input)
    print('out_conv2:',out_conv2.size())

    features = []

    # hook The function takes three arguments , These three parameters are transmitted by the system to hook Functional , You cannot modify these three parameters by yourself 
    def hook(module, input, output):
        features.append(output.clone().detach())

    net = model #LeNet()
    x = input
    #  Take out the corresponding layer of the network , Call... On this layer register_forward_hook Method . This method needs to pass in a hook Method ：
    handle = net.conv2.register_forward_hook(hook)
    # From here we can find hook You can even change the input and output ( But it will not affect the network forward The actual result of ), But here we simply put output Keep it .
    #  It should be noted that hook Functions should be deleted in time after use , To avoid increasing the running load every time .
    y,_ = net(x)
    print('net.conv2.register_forward_hook:',features[0].size())
    print('validation:',torch.allclose(features[0], out_conv2) )
    handle.remove()