Conocimiento de aprendizaje profundo 6: (modelo de compresión cuantitativa) ---- Módulo personalizado de pytorch, y comprenda el método de definición de red DoReFaNet a través de él.

Consulte al funcionario chino, para obtener más detalles, consulte: Cómo personalizar el módulo en PyTorch

1. Módulo personalizado

El módulo es la forma básica en que pytorch organiza las redes neuronales. El módulo contiene parámetros de modelo y lógica de cálculo. La función lleva la función real y define la lógica de cálculo hacia adelante y hacia atrás.
A continuación, se toma la estructura de red MLP más simple como ejemplo para presentar cómo implementar una estructura de red personalizada. El código completo se puede encontrar en el repositorio .

1.1 Función

La función es la clase principal del mecanismo de derivación automática de pytorch. La función es sin parámetros o apátrida, solo se encarga de recibir la entrada y devolver la salida correspondiente, para revertir recibe el gradiente correspondiente de la salida y devuelve el gradiente correspondiente de la entrada.

Aquí solo nos enfocamos en cómo personalizar la función. Consulte el código fuente para conocer la definición de Función . Aquí hay un fragmento de código simplificado:

class Function(object):
    def forward(self, *input):
        raise NotImplementedError
 
    def backward(self, *grad_output):
        raise NotImplementedError

La entrada y salida de avance y retroceso son objetos Tensor.

El objeto Function es invocable, es decir, puede ser llamado por (). La entrada y salida de la llamada son objetos variables. El siguiente ejemplo de código muestra cómo implementar una función de activación de ReLU y llamarla:

import torch
from torch.autograd import Function
 
class ReLUF(Function)：
    def forward(self, input):
        self.save_for_backward(input)
 
        output = input.clamp(min=0)
        return output
 
    def backward(self, output_grad):
        input = self.to_save[0]
 
        input_grad = output_grad.clone()
        input_grad[input < 0] = 0
        return input_grad
 
## Test
if __name__ == "__main__":
      from torch.autograd import Variable
 
      torch.manual_seed(1111)  
      a = torch.randn(2, 3)
 
      va = Variable(a, requires_grad=True)
      vb = ReLUF()(va)
      print va.data, vb.data
 
      vb.backward(torch.ones(va.size()))
      print vb.grad.data, va.grad.data

Si la entrada de avance es necesaria en el retroceso, debe guardar explícitamente la entrada requerida en el avance. En el código anterior, forward usa la función self.save_for_backward para guardar temporalmente la entrada, y usa Saved_tensors (objeto de tupla de Python) para sacarla al revés.

Obviamente, la entrada hacia adelante debe corresponder a la entrada hacia atrás; al mismo tiempo, la salida hacia adelante debe coincidir con la entrada hacia atrás.
Nota:
Dado que Function puede necesitar almacenar temporalmente el tensor de entrada, se recomienda no reutilizar el objeto Function para evitar el problema de la liberación anticipada de la memoria.
Como se muestra en el código de muestra , cada llamada a reenviar regenera un objeto ReLUF y no se puede generar durante la inicialización y las llamadas repetidas en reenvío. (Lo que significa que el código de muestra es correcto, es que cada fc genera un nuevo objeto. Porque cada objeto administra su propio peso y entrada )

2 módulos

Al igual que la función, el objeto Módulo también es invocable y la entrada y la salida también son variables. La diferencia es que Module puede tener parámetros. El módulo consta de dos partes principales: parámetros y lógica de cálculo (llamada de función). Dado que la función de activación de ReLU no tiene parámetros, aquí tomamos la capa completamente conectada más básica como ejemplo para ilustrar cómo personalizar el Módulo.
La lógica de funcionamiento de la capa completamente conectada se define de la siguiente manera Función:

import torch
from torch.autograd import Function
 
class LinearF(Function):
 
     def forward(self, input, weight, bias=None):
         self.save_for_backward(input, weight, bias)
 
         output = torch.mm(input, weight.t())
         if bias is not None:
             output += bias.unsqueeze(0).expand_as(output)
 
         return output
 
     def backward(self, grad_output):
         input, weight, bias = self.saved_tensors
 
         grad_input = grad_weight = grad_bias = None
         if self.needs_input_grad[0]:
             grad_input = torch.mm(grad_output, weight)
         if self.needs_input_grad[1]:
             grad_weight = torch.mm(grad_output.t(), input)
         if bias is not None and self.needs_input_grad[2]:
             grad_bias = grad_output.sum(0).squeeze(0)
 
         if bias is not None:
             return grad_input, grad_weight, grad_bias
         else:
             return grad_input, grad_weight

need_input_grad es una tupla con un elemento bool, la longitud es la misma que el número de parámetros de avance y se usa para identificar si cada entrada es una entrada para calcular el gradiente; para las entradas que no requieren gradientes, se pueden reducir los cálculos innecesarios .

La función (aquí LinearF) define la lógica de cálculo básica. El módulo solo necesita asignar espacio de memoria para los parámetros durante la inicialización y pasar los parámetros al objeto de función correspondiente durante el cálculo. el código se muestra a continuación:

import torch
import torch.nn as nn
 
class Linear(nn.Module):
 
    def __init__(self, in_features, out_features, bias=True):
         super(Linear, self).__init__()
         self.in_features = in_features
         self.out_features = out_features
         self.weight = nn.Parameter(torch.Tensor(out_features, in_features))
         if bias:
             self.bias = nn.Parameter(torch.Tensor(out_features))
         else:
            self.register_parameter('bias', None)
 
    def forward(self, input):
         return LinearF()(input, self.weight, self.bias)

Cabe señalar que el parámetro es el espacio de memoria mantenido por el objeto tensor, pero el tensor debe empaquetarse como un objeto Parameter . El parámetro es una subclase especial de Variable , la única diferencia es que el require_grad predeterminado de Parameter es Verdadero. Varaible es la clase principal del mecanismo de derivación automática. No se presentará aquí, consulte el tutorial .

2.3 Red neuronal recurrente personalizada (RNN)

Referencia de código ejecutable : RNN
Entre ellos, Parámetros es una subclase de Variable, y es un mecanismo de derivación automática, por lo que cuando definimos la red de código, básicamente no necesitamos estar al revés , siempre que el avance esté definido.

3 Definir DoReFaNet

Dado que comprime pesos, solo necesitamos definir una clase de convolución, una clase de función de activación y luego obtener los pesos en él, cuantificar y definir la red de la siguiente manera:

class AlexNet_Q(nn.Module):
  def __init__(self, wbit, abit, num_classes=1000):
    super(AlexNet_Q, self).__init__()
    Conv2d = conv2d_Q_fn(w_bit=wbit)
    Linear = linear_Q_fn(w_bit=wbit)

    self.features = nn.Sequential(
      nn.Conv2d(3, 96, kernel_size=11, stride=4, padding=2),
      nn.BatchNorm2d(96),
      nn.ReLU(inplace=True),
      nn.MaxPool2d(kernel_size=3, stride=2),

      Conv2d(96, 256, kernel_size=5, padding=2),
      nn.BatchNorm2d(256),
      nn.ReLU(inplace=True),
      activation_quantize_fn(a_bit=abit),
      nn.MaxPool2d(kernel_size=3, stride=2),

      Conv2d(256, 384, kernel_size=3, padding=1),
      nn.ReLU(inplace=True),
      activation_quantize_fn(a_bit=abit),

      Conv2d(384, 384, kernel_size=3, padding=1),
      nn.ReLU(inplace=True),
      activation_quantize_fn(a_bit=abit),

      Conv2d(384, 256, kernel_size=3, padding=1),
      nn.ReLU(inplace=True),
      activation_quantize_fn(a_bit=abit),
      nn.MaxPool2d(kernel_size=3, stride=2),
    )
    self.classifier = nn.Sequential(
      Linear(256 * 6 * 6, 4096),
      nn.ReLU(inplace=True),
      activation_quantize_fn(a_bit=abit),

      Linear(4096, 4096),
      nn.ReLU(inplace=True),
      activation_quantize_fn(a_bit=abit),
      nn.Linear(4096, num_classes),
    )

    for m in self.modules():
      if isinstance(m, Conv2d) or isinstance(m, Linear):
        init.xavier_normal_(m.weight.data)

  def forward(self, x):
    x = self.features(x)
    x = x.view(x.size(0), 256 * 6 * 6)
    x = self.classifier(x)
    return 

• Los derechos de personalización se cuantifican de la siguiente manera:

class weight_quantize_fn(nn.Module):
  def __init__(self, w_bit):
    super(weight_quantize_fn, self).__init__()
    assert w_bit <= 8 or w_bit == 32
    self.w_bit = w_bit
    self.uniform_q = uniform_quantize(k=w_bit)

  def forward(self, x):
    if self.w_bit == 32:
      weight_q = x
    elif self.w_bit == 1:
      E = torch.mean(torch.abs(x)).detach()
      weight_q = self.uniform_q(x / E) * E
    else:
      weight = torch.tanh(x)
      weight = weight / 2 / torch.max(torch.abs(weight)) + 0.5
      weight_q = 2 * self.uniform_q(weight) - 1
    return weight_q

def conv2d_Q_fn(w_bit):
  class Conv2d_Q(nn.Conv2d):
    def __init__(self, in_channels, out_channels, kernel_size, stride=1,
                 padding=0, dilation=1, groups=1, bias=True):
      super(Conv2d_Q, self).__init__(in_channels, out_channels, kernel_size, stride,
                                     padding, dilation, groups, bias)
      self.w_bit = w_bit
      self.quantize_fn = weight_quantize_fn(w_bit=w_bit)

    def forward(self, input, order=None):
      weight_q = self.quantize_fn(self.weight)
      # print(np.unique(weight_q.detach().numpy()))
      return F.conv2d(input, weight_q, self.bias, self.stride,
                      self.padding, self.dilation, self.groups)

  return Conv2d_Q

Entre ellos, la función de cuantificación de bits uniform_quantize es:

def uniform_quantize(k):
  class qfn(torch.autograd.Function):

    @staticmethod
    def forward(ctx, input):
      if k == 32:
        out = input
      elif k == 1:
        out = torch.sign(input)
      else:
        n = float(2 ** k - 1)
        out = torch.round(input * n) / n
      return out

    @staticmethod
    def backward(ctx, grad_output):
      grad_input = grad_output.clone()
      return grad_input

  return qfn().apply

• Una vez activada la capa de activación, los bits de salida se cuantifican como:

class activation_quantize_fn(nn.Module):
  def __init__(self, a_bit):
    super(activation_quantize_fn, self).__init__()
    assert a_bit <= 8 or a_bit == 32
    self.a_bit = a_bit
    self.uniform_q = uniform_quantize(k=a_bit)

  def forward(self, x):
    if self.a_bit == 32:
      activation_q = x
    else:
      activation_q = self.uniform_q(torch.clamp(x, 0, 1))
      # print(np.unique(activation_q.detach().numpy()))
    return activation_q

4. Realice la compresión del canal del modelo, utilizando el código yolov3 como ejemplo:

La compresión del canal se produce principalmente a través de BN的aX+b=ouputla distribución del peso a de la disposición . Después del entrenamiento, se ordena el peso a, se obtiene el umbral de acuerdo con la relación de cultivo, y luego se excluyen los canales menores que el umbral y se retienen los mayores que el umbral.

• Obtenga el índice de la cantidad de capas de red, incluida la capa BN:

def parse_module_defs(module_defs):
    CBL_idx = []
    Conv_idx = []
    for i, module_def in enumerate(module_defs):
        # 具有bn层的，才进行保存索引，然后用来压缩训练。
        if module_def['type'] == 'convolutional':
            if module_def['batch_normalize'] == '1':
                CBL_idx.append(i)
            else:
                Conv_idx.append(i)
    #排除一些不要进行压缩的层，为了保证准确度。
    ignore_idx = set()
    for i, module_def in enumerate(module_defs):
        if module_def['type'] == 'shortcut':
            ignore_idx.add(i-1)
            identity_idx = (i + int(module_def['from']))
            if module_defs[identity_idx]['type'] == 'convolutional':
                ignore_idx.add(identity_idx)
            elif module_defs[identity_idx]['type'] == 'shortcut':
                ignore_idx.add(identity_idx - 1)

    ignore_idx.add(84)
    ignore_idx.add(96)
    #获取符合要求跟需要压缩的网络层索引。
    prune_idx = [idx for idx in CBL_idx if idx not in ignore_idx]

    return CBL_idx, Conv_idx, prune_idx

• Código para entrenamiento:

        for batch_i, (_, imgs, targets) in enumerate(dataloader):
            batches_done = len(dataloader) * epoch + batch_i

            imgs = imgs.to(device)
            targets = targets.to(device)

            loss, outputs = model(imgs, targets)

            optimizer.zero_grad()
            loss.backward()
            #进行模型权重的稀疏化训练，opt.s可以控制稀疏值，prune_idx是包含需要进行稀疏训练的
            #卷积层索引。
            BNOptimizer.updateBN(sr_flag, model.module_list, opt.s, prune_idx)

            optimizer.step()

• La función principal de la cual el entrenamiento escaso, Principio: El opt.svalor adicional para hacer que el peso sea más importante para obtener una lata 正梯度外加值. La función del código es la siguiente:

class BNOptimizer():

    @staticmethod
    def updateBN(sr_flag, module_list, s, prune_idx):
        if sr_flag:
            for idx in prune_idx:
                # Squential(Conv, BN, Lrelu)
                bn_module = module_list[idx][1]
                #torch.sign()是一个跳变函数，其值只有-1，0，-1.
                #这里的s指的是附加到权重梯度grad上的值，其可以用来加大或减少梯度值，
                #从而在更新权重的时候，权值大的可以变的越来越大。
                bn_module.weight.grad.data.add_(s * torch.sign(bn_module.weight.data))  # L1

• Una vez entrenado todo el modelo, los canales de red innecesarios se pueden eliminar mediante la proporción de poda establecida por usted mismo. La función es la siguiente:

percent = 0.85
def prune_and_eval(model, sorted_bn, percent=.0):
    model_copy = deepcopy(model)
    #根据裁剪比例进行获取停止的通道索引
    thre_index = int(len(sorted_bn) * percent)
    #任何根据这个停止的通道索引获取其对应的通道权重值，
    #然后作为刷选阈值。
    thre = sorted_bn[thre_index]
    print(f'Channels with Gamma value less than {thre:.4f} are pruned!')
    remain_num = 0
    for idx in prune_idx:
        bn_module = model_copy.module_list[idx][1]
        #通过阈值获取刷选通道
        mask = obtain_bn_mask(bn_module, thre)
        remain_num += int(mask.sum())
        #根据mask的0\1进行获取保留了通道的
        bn_module.weight.data.mul_(mask)
    mAP = eval_model(model_copy)[2].mean()

    print(f'Number of channels has been reduced from {len(sorted_bn)} to {remain_num}')
    print(f'Prune ratio: {1-remain_num/len(sorted_bn):.3f}')
    print(f'mAP of the pruned model is {mAP:.4f}')

    return thre