├── README.md └── peleenet.py /README.md: -------------------------------------------------------------------------------- 1 | # PeleeNet in PyTorch 2 | 3 | This repository contains the peleenet (in PyTorch) for "[Pelee: A Real-Time Object Detection System on Mobile Devices](https://arxiv.org/pdf/1804.06882.pdf)". 4 | modified from [in Caffe](https://github.com/Robert-JunWang/Pelee/blob/master/peleenet.py)". 5 | 6 | 7 | 8 | 9 | ## Contact 10 | pfw813@gmail.com 11 | 12 | Any discussions or concerns are welcomed! 13 | 14 | 15 | ### Citation 16 | If you find this paper useful in your research, please consider citing: 17 | 18 | ``` 19 | @article{wang2018pelee, 20 | title={Pelee: A Real-Time Object Detection System on Mobile Devices}, 21 | author={Wang, Robert J and Li, Xiang and Ao, Shuang and Ling, Charles X}, 22 | journal={arXiv preprint arXiv:1804.06882}, 23 | year={2018} 24 | } 25 | ``` 26 | -------------------------------------------------------------------------------- /peleenet.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import math 4 | import torch.nn.functional as F 5 | 6 | 7 | class Conv_bn_relu(nn.Module): 8 | def __init__(self, inp, oup, kernel_size=3, stride=1, pad=1,use_relu = True): 9 | super(Conv_bn_relu, self).__init__() 10 | self.use_relu = use_relu 11 | if self.use_relu: 12 | self.convs = nn.Sequential( 13 | nn.Conv2d(inp, oup, kernel_size, stride, pad, bias=False), 14 | nn.BatchNorm2d(oup), 15 | nn.ReLU(inplace=True), 16 | ) 17 | else: 18 | self.convs = nn.Sequential( 19 | nn.Conv2d(inp, oup, kernel_size, stride, pad, bias=False), 20 | nn.BatchNorm2d(oup), 21 | ) 22 | 23 | 24 | def forward(self, x): 25 | out = self.convs(x) 26 | return out 27 | 28 | 29 | class StemBlock(nn.Module): 30 | def __init__(self, inp=3,num_init_features=32): 31 | super(StemBlock, self).__init__() 32 | 33 | 34 | 35 | self.stem_1 = Conv_bn_relu(inp, num_init_features, 3, 2, 1) 36 | 37 | self.stem_2a = Conv_bn_relu(num_init_features,int(num_init_features/2),1,1,0) 38 | 39 | self.stem_2b = Conv_bn_relu(int(num_init_features/2), num_init_features, 3, 2, 1) 40 | 41 | self.stem_2p = nn.MaxPool2d(kernel_size=2,stride=2) 42 | 43 | self.stem_3 = Conv_bn_relu(num_init_features*2,num_init_features,1,1,0) 44 | 45 | 46 | 47 | def forward(self, x): 48 | stem_1_out = self.stem_1(x) 49 | 50 | stem_2a_out = self.stem_2a(stem_1_out) 51 | stem_2b_out = self.stem_2b(stem_2a_out) 52 | 53 | stem_2p_out = self.stem_2p(stem_1_out) 54 | 55 | out = self.stem_3(torch.cat((stem_2b_out,stem_2p_out),1)) 56 | 57 | return out 58 | 59 | 60 | class DenseBlock(nn.Module): 61 | def __init__(self, inp,inter_channel,growth_rate): 62 | super(DenseBlock, self).__init__() 63 | 64 | self.cb1_a = Conv_bn_relu(inp,inter_channel,1,1,0) 65 | self.cb1_b = Conv_bn_relu(inter_channel,growth_rate,3,1,1) 66 | 67 | self.cb2_a = Conv_bn_relu(inp,inter_channel,1,1,0) 68 | self.cb2_b = Conv_bn_relu(inter_channel,growth_rate,3,1,1) 69 | self.cb2_c = Conv_bn_relu(growth_rate,growth_rate,3,1,1) 70 | 71 | 72 | def forward(self, x): 73 | cb1_a_out = self.cb1_a(x) 74 | cb1_b_out = self.cb1_b(cb1_a_out) 75 | 76 | cb2_a_out = self.cb2_a(x) 77 | cb2_b_out = self.cb2_b(cb2_a_out) 78 | cb2_c_out = self.cb2_c(cb2_b_out) 79 | 80 | out = torch.cat((x,cb1_b_out,cb2_c_out),1) 81 | 82 | return out 83 | 84 | 85 | class TransitionBlock(nn.Module): 86 | def __init__(self, inp, oup,with_pooling= True): 87 | super(TransitionBlock, self).__init__() 88 | if with_pooling: 89 | self.tb = nn.Sequential(Conv_bn_relu(inp,oup,1,1,0), 90 | nn.AvgPool2d(kernel_size=2,stride=2)) 91 | else: 92 | self.tb = Conv_bn_relu(inp,oup,1,1,0) 93 | 94 | def forward(self, x): 95 | out = self.tb(x) 96 | return out 97 | 98 | 99 | class PeleeNet(nn.Module): 100 | def __init__(self,num_classes=1000, num_init_features=32,growthRate=32, nDenseBlocks = [3,4,8,6], bottleneck_width=[1,2,4,4]): 101 | super(PeleeNet, self).__init__() 102 | 103 | 104 | self.stage = nn.Sequential() 105 | self.num_classes = num_classes 106 | self.num_init_features = num_init_features 107 | 108 | inter_channel =list() 109 | total_filter =list() 110 | dense_inp = list() 111 | 112 | self.half_growth_rate = int(growthRate / 2) 113 | 114 | # building stemblock 115 | self.stage.add_module('stage_0', StemBlock(3,num_init_features)) 116 | 117 | # 118 | for i, b_w in enumerate(bottleneck_width): 119 | 120 | inter_channel.append(int(self.half_growth_rate * b_w / 4) * 4) 121 | 122 | if i == 0: 123 | total_filter.append(num_init_features + growthRate * nDenseBlocks[i]) 124 | dense_inp.append(self.num_init_features) 125 | else: 126 | total_filter.append(total_filter[i-1] + growthRate * nDenseBlocks[i]) 127 | dense_inp.append(total_filter[i-1]) 128 | 129 | if i == len(nDenseBlocks)-1: 130 | with_pooling = False 131 | else: 132 | with_pooling = True 133 | 134 | # building middle stageblock 135 | self.stage.add_module('stage_{}'.format(i+1),self._make_dense_transition(dense_inp[i], total_filter[i], 136 | inter_channel[i],nDenseBlocks[i],with_pooling=with_pooling)) 137 | 138 | 139 | # building classifier 140 | self.classifier = nn.Sequential( 141 | nn.Dropout(), 142 | nn.Linear(total_filter[len(nDenseBlocks)-1], self.num_classes) 143 | ) 144 | 145 | self._initialize_weights() 146 | 147 | 148 | def _make_dense_transition(self, dense_inp,total_filter, inter_channel, ndenseblocks,with_pooling= True): 149 | layers = [] 150 | 151 | for i in range(ndenseblocks): 152 | layers.append(DenseBlock(dense_inp, inter_channel,self.half_growth_rate)) 153 | dense_inp += self.half_growth_rate * 2 154 | 155 | #Transition Layer without Compression 156 | layers.append(TransitionBlock(dense_inp,total_filter,with_pooling)) 157 | 158 | return nn.Sequential(*layers) 159 | 160 | def forward(self, x): 161 | 162 | x = self.stage(x) 163 | 164 | # global average pooling layer 165 | x = F.avg_pool2d(x,kernel_size=7) 166 | x = x.view(x.size(0), -1) 167 | x = self.classifier(x) 168 | out = F.log_softmax(x,dim=1) 169 | 170 | return out 171 | 172 | def _initialize_weights(self): 173 | for m in self.modules(): 174 | if isinstance(m, nn.Conv2d): 175 | n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels 176 | m.weight.data.normal_(0, math.sqrt(2. / n)) 177 | if m.bias is not None: 178 | m.bias.data.zero_() 179 | elif isinstance(m, nn.BatchNorm2d): 180 | m.weight.data.fill_(1) 181 | m.bias.data.zero_() 182 | elif isinstance(m, nn.Linear): 183 | n = m.weight.size(1) 184 | m.weight.data.normal_(0, 0.01) 185 | m.bias.data.zero_() 186 | 187 | 188 | 189 | if __name__ == '__main__': 190 | p = PeleeNet(num_classes=1000) 191 | input = torch.autograd.Variable(torch.ones(1, 3, 224, 224)) 192 | output = p(input) 193 | 194 | print(output.size()) 195 | 196 | # torch.save(p.state_dict(), 'peleenet.pth.tar') 197 | 198 | 199 | 200 | 201 | 202 | 203 | 204 | 205 | 206 | 207 | 208 | 209 | 210 | 211 | 212 | 213 | --------------------------------------------------------------------------------