ModulatedConv2d

  2. class ModulatedConv2d(nn.Module):
  3.     def __init__(self, fin, fout, kernel_size, padding_type='reflect', upsample=False, downsample=False, latent_dim=256, normalize_mlp=False):
  4.         super(ModulatedConv2d, self).__init__()
  5.         self.in_channels = fin
  6.         self.out_channels = fout
  7.         self.kernel_size = kernel_size
  8.         self.upsample = upsample
  9.         self.downsample = downsample
  10.         padding_size = kernel_size // 2
  11.         if kernel_size == 1:
  12.             self.demudulate = False
  13.         else:
  14.             self.demudulate = True
  16.         self.weight = nn.Parameter(torch.Tensor(fout, fin, kernel_size, kernel_size))
  17.         self.bias = nn.Parameter(torch.Tensor(1, fout, 1, 1))
  18.         self.conv = F.conv2d
  20.         if normalize_mlp:
  21.             self.mlp_class_std = nn.Sequential(EqualLinear(latent_dim, fin), PixelNorm())
  22.         else:
  23.             self.mlp_class_std = EqualLinear(latent_dim, fin)
  25.         self.blur = Blur(fout)
  27.         if padding_type == 'reflect':
  28.             self.padding = nn.ReflectionPad2d(padding_size)
  29.         else:
  30.             self.padding = nn.ZeroPad2d(padding_size)
  32.         if self.upsample:
  33.             self.upsampler = nn.Upsample(scale_factor=2, mode='nearest')
  35.         if self.downsample:
  36.             self.downsampler = nn.AvgPool2d(2)
  38.         self.weight.data.normal_()
  39.         self.bias.data.zero_()
  41.     def forward(self, input, latent):
  42.         fan_in = self.weight.data.size(1) * self.weight.data[0][0].numel()
  43.         weight = self.weight * sqrt(2 / fan_in)
  44.         weight = weight.view(1, self.out_channels, self.in_channels, self.kernel_size, self.kernel_size)
  46.         s = 1 + self.mlp_class_std(latent).view(-1, 1, self.in_channels, 1, 1)
  47.         weight = s * weight
  48.         if self.demudulate:
  49.             d = torch.rsqrt((weight ** 2).sum(4).sum(3).sum(2) + 1e-5).view(-1, self.out_channels, 1, 1, 1)
  50.             weight = (d * weight).view(-1, self.in_channels, self.kernel_size, self.kernel_size)
  51.         else:
  52.             weight = weight.view(-1, self.in_channels, self.kernel_size, self.kernel_size)
  54.         if self.upsample:
  55.             input = self.upsampler(input)
  57.         if self.downsample:
  58.             input = self.blur(input)
  60.         b,_,h,w = input.shape
  61.         input = input.view(1,-1,h,w)
  62.         input = self.padding(input)
  63.         out = self.conv(input, weight, groups=b).view(b, self.out_channels, h, w) + self.bias
  65.         if self.downsample:
  66.             out = self.downsampler(out)
  68.         if self.upsample:
  69.             out = self.blur(out)
  71.         return out
  1. from torch.autograd import Function
  3. class BlurFunctionBackward(Function):
  4.     @staticmethod
  5.     def forward(ctx, grad_output, kernel, kernel_flip):
  6.         ctx.save_for_backward(kernel, kernel_flip)
  8.         grad_input = F.conv2d(
  9.             grad_output, kernel_flip, padding=1, groups=grad_output.shape[1]
  10.         )
  12.         return grad_input
  14.     @staticmethod
  15.     def backward(ctx, gradgrad_output):
  16.         kernel, kernel_flip = ctx.saved_tensors
  18.         grad_input = F.conv2d(
  19.             gradgrad_output, kernel, padding=1, groups=gradgrad_output.shape[1]
  20.         )
  22.         return grad_input, None, None
  24. class BlurFunction(Function):
  25.     @staticmethod
  26.     def forward(ctx, input, kernel, kernel_flip):
  27.         ctx.save_for_backward(kernel, kernel_flip)
  29.         output = F.conv2d(input, kernel, padding=1, groups=input.shape[1])
  31.         return output
  33.     @staticmethod
  34.     def backward(ctx, grad_output):
  35.         kernel, kernel_flip = ctx.saved_tensors
  37.         grad_input = BlurFunctionBackward.apply(grad_output, kernel, kernel_flip)
  39.         return grad_input, None, None
  41. blur = BlurFunction.apply
  43. class Blur(nn.Module):
  44.     def __init__(self, channel):
  45.         super().__init__()
  47.         weight = torch.tensor([[1, 2, 1], [2, 4, 2], [1, 2, 1]], dtype=torch.float32)
  48.         weight = weight.view(1, 1, 3, 3)
  49.         weight = weight / weight.sum()
  50.         weight_flip = torch.flip(weight, [2, 3])
  52.         self.register_buffer('weight', weight.repeat(channel, 1, 1, 1))
  53.         self.register_buffer('weight_flip', weight_flip.repeat(channel, 1, 1, 1))
  55.     def forward(self, input):
  56.         return blur(input, self.weight, self.weight_flip)

0. 本章内容

在本次，我们将学习如何自定义一个torch.autograd.Function，下面是本次的主要内容

  1. 对Function的直观理解
  2. Function与Module的差异与应用场景
  3. 写一个简单的ReLU Function

    1.对Function的直观理解

  • 在之前的介绍中，我们知道，Pytorch是利用Variable与Function来构建计算图的。回顾下Variable，Variable就像是计算图中的节点，保存计算结果（包括前向传播的激活值，反向传播的梯度），而Function就像计算图中的边，实现Variable的计算，并输出新的Variable
  • Function简单说就是对Variable的运算，如加减乘除，relu，pool等
  • 但它不仅仅是简单的运算。与普通Python或者numpy的运算不同，Function是针对计算图，需要计算反向传播的梯度。因此他不仅需要进行该运算（forward过程），还需要保留前向传播的输入（为计算梯度），并支持反向传播计算梯度。如果有做过公开课cs231的作业，记得里面的每个运算都定义了forward，backward，并通过保存cache来进行反向传播。这两者是类似的。
  • 在之前Variable的学习中，我们知道进行一次运算后，输出的Variable对应的creator就是其运行的计算，如y = relu(x), y.creator，就是relu这个Function
  • 我们可以对Function进行拓展，使其满足我们自己的需要，而拓展就需要自定义Function的forward运算，已经对应的backward运算，同时在forward中需要通过保存输入值用于backward

  • 总结，Function与Variable构成了pytorch的自动求导机制，它定义的是各个Variable之间的计算关系

    2. Function与Module的差异与应用场景

    Function与Module都可以对pytorch进行自定义拓展，使其满足网络的需求，但这两者还是有十分重要的不同：

  • Function一般只定义一个操作，因为其无法保存参数，因此适用于激活函数、pooling等操作；Module是保存了参数，因此适合于定义一层，如线性层，卷积层，也适用于定义一个网络

  • Function需要定义三个方法：init, forward, backward（需要自己写求导公式）；Module：只需定义init和forward，而backward的计算由自动求导机制构成
  • 可以不严谨的认为，Module是由一系列Function组成，因此其在forward的过程中，Function和Variable组成了计算图，在backward时，只需调用Function的backward就得到结果，因此Module不需要再定义backward。
  • Module不仅包括了Function，还包括了对应的参数，以及其他函数与变量，这是Function所不具备的

    3. 一个ReLU Function

  1. 首先我们定义一个继承Function的ReLU类
  2. 然后我们来看Variable在进行运算时，其creator是否是对应的Function
  3. 最后我们为方便使用这个ReLU类，将其wrap成一个函数，方便调用，不必每次显式都创建一个新对象

3.1 定义一个ReLU类

  1. import torch
  2. from torch.autograd import Variable
  3. class MyReLU(torch.autograd.Function):
  4.     def forward(self, input_):
  5.         # 在forward中，需要定义MyReLU这个运算的forward计算过程
  6.         # 同时可以保存任何在后向传播中需要使用的变量值
  7.         self.save_for_backward(input_)         # 将输入保存起来，在backward时使用
  8.         output = input_.clamp(min=0)               # relu就是截断负数，让所有负数等于0
  9.         return output
  10.     def backward(self, grad_output):
  11.         # 根据BP算法的推导（链式法则），dloss / dx = (dloss / doutput) * (doutput / dx)
  12.         # dloss / doutput就是输入的参数grad_output、
  13.         # 因此只需求relu的导数，在乘以grad_outpu    
  14.         input_, = self.saved_tensors
  15.         grad_input = grad_output.clone()
  16.         grad_input[input < 0] = 0                # 上诉计算的结果就是左式。即ReLU在反向传播中可以看做一个通道选择函数，所有未达到阈值（激活值<0）的单元的梯度都为0
  17.         return grad_input

3.2 验证Variable与Function的关系

  1. from torch.autograd import Variable
  2. input_ = Variable(torch.randn(1))
  3. relu = MyReLU()
  4. output_ = relu(input_)
  5. # 这个relu对象，就是output_.creator，即这个relu对象将output与input连接起来，形成一个计算图
  6. print relu
  7. print output_.creator

输出：

  1. <__main__.MyReLU object at 0x7fd0b2d08b30>
  2. <__main__.MyReLU object at 0x7fd0b2d08b30>

可见，Function连接了Variable与Variable并实现不同计算
3.3 Wrap一个ReLU函数
可以直接把刚才自定义的ReLU类封装成一个函数，方便直接调用

  1. def relu(input_):
  2.     # MyReLU()是创建一个MyReLU对象，
  3.     # Function类利用了Python __call__操作，使得可以直接使用对象调用__call__制定的方法
  4.     # __call__指定的方法是forward，因此下面这句MyReLU（）（input_）相当于
  5.     # return MyReLU().forward(input_)
  6.     return MyReLU()(input_)
  7. input_ = Variable(torch.linspace(-3, 3, steps=5))
  8. print input_
  9. print relu(input_)

输出：

  1. Variable containing:
  2. -3.0000
  3. -1.5000
  4.  0.0000
  5.  1.5000
  6.  3.0000
  7. [torch.FloatTensor of size 5]
  8. Variable containing:
  9.  0.0000
  10.  0.0000
  11.  0.0000
  12.  1.5000
  13.  3.0000
  14. [torch.FloatTensor of size 5]