Pytorch 抽取vgg各層并進(jìn)行定制化處理的方法

發(fā)布時間：2020-09-16 21:04:01 來源：腳本之家閱讀：120 作者：xiaoxifei 欄目：開發(fā)技術(shù)

工作中有時候需要對vgg進(jìn)行定制化處理，比如有些時候需要借助于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ìn)行層的提??；也可以根據(jù)定制在classifier上進(jìn)行提取。

下面是我的一個提取前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各層并進(jìn)行定制化處理的方法就是小編分享給大家的全部內(nèi)容了，希望能給大家一個參考，也希望大家多多支持億速云。

向AI問一下細(xì)節(jié)

Pytorch 抽取vgg各層并進(jìn)行定制化處理的方法

猜你喜歡

最新資訊

相關(guān)推薦

相關(guān)標(biāo)簽