Pytorch 抽取vgg各层并进行定制化处理的方法
工作中有时候需要对vgg进行定制化处理,比如有些时候需要借助于vgg的层结构,但是需要使用的是2channels输入,等等需求,这时候可以使用vgg的原始结构用class重写一遍,但是这样的方式比较慢,并且容易出错,下面给出一种比较简单的方式
defdefine_vgg(vgg,input_channels,endlayer,use_maxpool=False): vgg_ad=copy.deepcopy(vgg) model=nn.Sequential() i=0 forlayerinlist(vgg_ad.features): ifi>endlayer: break ifisinstance(layer,nn.Conv2d)andiis0: 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) ifisinstance(layer,nn.Conv2d): name="conv_"+str(i) model.add_module(name,layer) ifisinstance(layer,nn.ReLU): name="leakyrelu_"+str(i) layer=nn.LeakyReLU(inplace=True) model.add_module(name,layer) ifisinstance(layer,nn.MaxPool2d): name="pool_"+str(i) ifuse_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 returnmodel
函数输入项中的vgg是直接使用的importtorchvision.models.vgg16传入的是vgg16非预训练版本。end_layer是需要提取的层数,这里使用了vgg.features是指仅仅在vgg.features上进行层的提取;也可以根据定制在classifier上进行提取。
下面是我的一个提取前7层的示例,可以使用pyCharmevaluate上面函数返回的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各层并进行定制化处理的方法就是小编分享给大家的全部内容了,希望能给大家一个参考,也希望大家多多支持毛票票。