├── README.md ├── SKNet.py ├── __init__.py └── timg.jpg /README.md: -------------------------------------------------------------------------------- 1 | # SKNet_pytorch 2 | This is my reimplementation the [`Selective Kernel Networks(SKNet)(CVPR2019)`](https://arxiv.org/pdf/1903.06586.pdf "the CVPR2019 paper") by pytorch
3 | The original implementation can get from [`here`](https://github.com/implus/SKNet "the implementation by caffe") 4 | 5 | ## Requiresments 6 | Python 3.6, [`Pytorch 1.0`](https://pytorch.org/?source=Google&medium=PaidSearch&utm_campaign=%2A%2ALP+-+TM+-+General+-+HV+-+HK&utm_adgroup=PyTorch+1.0&utm_keyword=pytorch%201.0&utm_offering=AI&utm_Product=PyTorch&gclid=Cj0KCQjws5HlBRDIARIsAOomqA2cDM3B1mvkQafcWtMeYjTuzyyL9QErEf-zXvoHpy0BM6hWR4cn6t4aAm0NEALw_wcB "Pytorch1.0") 7 | -------------------------------------------------------------------------------- /SKNet.py: -------------------------------------------------------------------------------- 1 | import torch.nn as nn 2 | from functools import reduce 3 | class SKConv(nn.Module): 4 | def __init__(self,in_channels,out_channels,stride=1,M=2,r=16,L=32): 5 | super(SKConv,self).__init__() 6 | d=max(in_channels//r,L) 7 | self.M=M 8 | self.out_channels=out_channels 9 | self.conv=nn.ModuleList() 10 | for i in range(M): 11 | self.conv.append(nn.Sequential(nn.Conv2d(in_channels,out_channels,3,stride,padding=1+i,dilation=1+i,groups=32,bias=False), 12 | nn.BatchNorm2d(out_channels), 13 | nn.ReLU(inplace=True))) 14 | self.global_pool=nn.AdaptiveAvgPool2d(1) 15 | self.fc1=nn.Sequential(nn.Conv2d(out_channels,d,1,bias=False), 16 | nn.BatchNorm2d(d), 17 | nn.ReLU(inplace=True)) 18 | self.fc2=nn.Conv2d(d,out_channels*M,1,1,bias=False) 19 | self.softmax=nn.Softmax(dim=1) 20 | def forward(self, input): 21 | batch_size=input.size(0) 22 | output=[] 23 | #the part of split 24 | for i,conv in enumerate(self.conv): 25 | #print(i,conv(input).size()) 26 | output.append(conv(input)) 27 | #the part of fusion 28 | U=reduce(lambda x,y:x+y,output) 29 | s=self.global_pool(U) 30 | z=self.fc1(s) 31 | a_b=self.fc2(z) 32 | a_b=a_b.reshape(batch_size,self.M,self.out_channels,-1) 33 | a_b=self.softmax(a_b) 34 | #the part of selection 35 | a_b=list(a_b.chunk(self.M,dim=1))#split to a and b 36 | a_b=list(map(lambda x:x.reshape(batch_size,self.out_channels,1,1),a_b)) 37 | V=list(map(lambda x,y:x*y,output,a_b)) 38 | V=reduce(lambda x,y:x+y,V) 39 | return V 40 | class SKBlock(nn.Module): 41 | expansion=2 42 | def __init__(self,inplanes,planes,stride=1,downsample=None): 43 | super(SKBlock,self).__init__() 44 | self.conv1=nn.Sequential(nn.Conv2d(inplanes,planes,1,1,0,bias=False), 45 | nn.BatchNorm2d(planes), 46 | nn.ReLU(inplace=True)) 47 | self.conv2=SKConv(planes,planes,stride) 48 | self.conv3=nn.Sequential(nn.Conv2d(planes,planes*self.expansion,1,1,0,bias=False), 49 | nn.BatchNorm2d(planes*self.expansion)) 50 | self.relu=nn.ReLU(inplace=True) 51 | self.downsample=downsample 52 | def forward(self, input): 53 | shortcut=input 54 | output=self.conv1(input) 55 | output=self.conv2(output) 56 | output=self.conv3(output) 57 | if self.downsample is not None: 58 | shortcut=self.downsample(input) 59 | output+=shortcut 60 | return self.relu(output) 61 | class SKNet(nn.Module): 62 | def __init__(self,nums_class=1000,block=SKBlock,nums_block_list=[3, 4, 6, 3]): 63 | super(SKNet,self).__init__() 64 | self.inplanes=64 65 | self.conv=nn.Sequential(nn.Conv2d(3,64,7,2,3,bias=False), 66 | nn.BatchNorm2d(64), 67 | nn.ReLU(inplace=True)) 68 | self.maxpool=nn.MaxPool2d(3,2,1) 69 | self.layer1=self._make_layer(block,128,nums_block_list[0],stride=1) 70 | self.layer2=self._make_layer(block,256,nums_block_list[1],stride=2) 71 | self.layer3=self._make_layer(block,512,nums_block_list[2],stride=2) 72 | self.layer4=self._make_layer(block,1024,nums_block_list[3],stride=2) 73 | self.avgpool=nn.AdaptiveAvgPool2d(1) 74 | self.fc=nn.Linear(1024*block.expansion,nums_class) 75 | self.softmax=nn.Softmax(-1) 76 | def forward(self, input): 77 | output=self.conv(input) 78 | output=self.maxpool(output) 79 | output=self.layer1(output) 80 | output=self.layer2(output) 81 | output=self.layer3(output) 82 | output=self.layer4(output) 83 | output=self.avgpool(output) 84 | output=output.squeeze(-1).squeeze(-1) 85 | output=self.fc(output) 86 | output=self.softmax(output) 87 | return output 88 | def _make_layer(self,block,planes,nums_block,stride=1): 89 | downsample=None 90 | if stride!=1 or self.inplanes!=planes*block.expansion: 91 | downsample=nn.Sequential(nn.Conv2d(self.inplanes,planes*block.expansion,1,stride,bias=False), 92 | nn.BatchNorm2d(planes*block.expansion)) 93 | layers=[] 94 | layers.append(block(self.inplanes,planes,stride,downsample)) 95 | self.inplanes=planes*block.expansion 96 | for _ in range(1,nums_block): 97 | layers.append(block(self.inplanes,planes)) 98 | return nn.Sequential(*layers) 99 | def SKNet50(nums_class=1000): 100 | return SKNet(nums_class,SKBlock,[3, 4, 6, 3]) 101 | def SKNet101(nums_class=1000): 102 | return SKNet(nums_class,SKBlock,[3, 4, 23, 3]) 103 | if __name__=='__main__': 104 | from PIL import Image 105 | from torchvision import transforms 106 | from torch.autograd import Variable 107 | import torch 108 | img=Image.open('timg.jpg').convert('RGB') 109 | img=transforms.ToTensor()(img) 110 | img=Variable(img).cuda() 111 | img=torch.stack([img,img]) 112 | #img=img.unsqueeze(0) 113 | temp=SKNet50().cuda() 114 | pred=temp(img) 115 | print(pred) 116 | 117 | -------------------------------------------------------------------------------- /__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ResearchingDexter/SKNet_pytorch/1e1ce1a221414f249ac4a898f99ec9e3ffaf09fd/__init__.py -------------------------------------------------------------------------------- /timg.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ResearchingDexter/SKNet_pytorch/1e1ce1a221414f249ac4a898f99ec9e3ffaf09fd/timg.jpg --------------------------------------------------------------------------------