快捷導(dǎo)航

Pytorch 抽取vgg各層并進行定制化處理的方法

更新時間：2019年08月20日 09:12:14 作者：xiaoxifei

今天小編就為大家分享一篇Pytorch 抽取vgg各層并進行定制化處理的方法，具有很好的參考價值，希望對大家有所幫助。一起跟隨小編過來看看吧

工作中有時候需要對vgg進行定制化處理，比如有些時候需要借助于vgg的層結(jié)構(gòu)，但是需要使用的是2 channels輸入，等等需求，這時候可以使用vgg的原始結(jié)構(gòu)用class重寫一遍，但是這樣的方式比較慢，并且容易出錯，下面給出一種比較簡單的方式

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

函數(shù)輸入項中的vgg 是直接使用的import torchvision.models.vgg16 傳入的是vgg16 非預(yù)訓(xùn)練版本。end_layer 是需要提取的層數(shù)，這里使用了vgg.features 是指僅僅在vgg.features 上進行層的提??；也可以根據(jù)定制在classifier上進行提取。

下面是我的一個提取前7層的示例，可以使用pyCharm evaluate 上面函數(shù)返回的model，可以看到這個示例的情況，這里我的定制條件是輸入通道為2 ，需要提取前7層，并且將ReLu更換為LeakyRelu。

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))
)

以上這篇Pytorch 抽取vgg各層并進行定制化處理的方法就是小編分享給大家的全部內(nèi)容了，希望能給大家一個參考，也希望大家多多支持腳本之家。

您可能感興趣的文章: