Pytorch 抽取vgg各层并进行定制化处理的方法

def define_vgg(vgg,input_channels,endlayer,use_maxpool=False): 
  vgg_ad = copy.deepcopy(vgg)
  model = nn.Sequential()
  i = 0
  for layer in list(vgg_ad.features):
    if i > endlayer:
      break
    if isinstance(layer, nn.Conv2d) and i is 0:
      name = "conv_" + str(i)
      layer = nn.Conv2d(input_channels,
               layer.out_channels,
               layer.kernel_size,
               stride = layer.stride,
               padding=layer.padding)
      model.add_module(name, layer)
    if isinstance(layer, nn.Conv2d):
      name = "conv_" + str(i)
      model.add_module(name, layer)
 
    if isinstance(layer, nn.ReLU):
      name = "leakyrelu_" + str(i)
      layer = nn.LeakyReLU(inplace=True) 
      model.add_module(name, layer)
 
    if isinstance(layer, nn.MaxPool2d):
      name = "pool_" + str(i)
      if use_maxpool:
        model.add_module(name, layer)
      else:
        avgpool = nn.AvgPool2d(kernel_size=layer.kernel_size, stride=layer.stride, padding=layer.padding)
        model.add_module(name, avgpool)
    i += 1
  return model

Sequential(
 (conv_0): Conv2d(2, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
 (leakyrelu_1): LeakyReLU(negative_slope=0.01, inplace)
 (conv_2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
 (leakyrelu_3): LeakyReLU(negative_slope=0.01, inplace)
 (pool_4): AvgPool2d(kernel_size=2, stride=2, padding=0)
 (conv_5): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
 (leakyrelu_6): LeakyReLU(negative_slope=0.01, inplace)
 (conv_7): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
)