forked from TensorLayer/tensorlayer3
225 lines
6.3 KiB
Python
225 lines
6.3 KiB
Python
from __future__ import absolute_import
|
|
from __future__ import division
|
|
from __future__ import print_function
|
|
from config import cfg
|
|
from network import Network
|
|
import tensorlayer as tl
|
|
from tensorlayer.layers import (BatchNorm, Conv2d, Dense, MaxPool2d , MeanPool2d)
|
|
from tensorlayer.layers import Module, SequentialLayer
|
|
import os
|
|
os.environ['TL_BACKEND'] = 'tensorflow'
|
|
import math
|
|
import torch
|
|
|
|
__all__ = ['ResNet', 'resnet101',]
|
|
|
|
|
|
def conv3x3(in_planes, out_planes, stride=1):
|
|
"3x3 convolution with padding"
|
|
return Conv2d(out_planes,(3,3),(1,1),in_channels=in_planes)
|
|
|
|
|
|
class BasicBlock(Module):
|
|
expansion = 1
|
|
|
|
def __init__(self, inplanes, planes, stride=1, downsample=None):
|
|
super(BasicBlock, self).__init__()
|
|
self.conv1 = conv3x3(inplanes, planes, stride)
|
|
self.bn1 = BatchNorm(planes)
|
|
self.conv2 = conv3x3(planes, planes)
|
|
self.bn2 = BatchNorm(planes)
|
|
self.downsample = downsample
|
|
self.stride = stride
|
|
|
|
def forward(self, x):
|
|
residual = x
|
|
|
|
out = self.conv1(x)
|
|
out = self.bn1(out,act='relu')
|
|
|
|
out = self.conv2(out)
|
|
out = self.bn2(out,act='relu')
|
|
|
|
if self.downsample is not None:
|
|
residual = self.downsample(x)
|
|
|
|
out += residual
|
|
out = tl.relu(out)
|
|
return out
|
|
|
|
|
|
class Bottleneck(Module):
|
|
expansion = 4
|
|
|
|
def __init__(self, inplanes, planes, stride=1, downsample=None):
|
|
super(Bottleneck, self).__init__()
|
|
|
|
self.conv1 = Conv2d(planes,(1,1), (stride,stride),in_channels=inplanes) # change
|
|
self.bn1 = BatchNorm(planes)
|
|
|
|
self.conv2 = Conv2d(planes,(3,3),(1,1),in_channels=planes)
|
|
self.bn2 = BatchNorm(planes)
|
|
|
|
self.conv3 = Conv2d(planes * 4,(1,1),(1,1),in_channels=planes)
|
|
self.bn3 = BatchNorm(planes * 4)
|
|
self.downsample = downsample
|
|
self.stride = stride
|
|
|
|
def forward(self, x):
|
|
residual = x
|
|
|
|
out = self.conv1(x)
|
|
out = self.bn1(out,act='relu')
|
|
|
|
|
|
out = self.conv2(out)
|
|
out = self.bn2(out,act='relu')
|
|
|
|
out = self.conv3(out)
|
|
out = self.bn3(out)
|
|
|
|
if self.downsample is not None:
|
|
residual = self.downsample(x)
|
|
|
|
out += residual
|
|
out = tl.relu(out)
|
|
|
|
return out
|
|
|
|
|
|
class ResNet(Module):
|
|
def __init__(self, block, layers, num_classes=1000):
|
|
self.inplanes = 64
|
|
super(ResNet, self).__init__()
|
|
self.conv1 = Conv2d(64,(7,7),(2,2),padding='valid',in_channels=3)
|
|
self.bn1 = BatchNorm(64)
|
|
self.maxpool = MaxPool2d((3,3),(2,2)) # change
|
|
self.layer1 = self._make_layer(block, 64, layers[0])
|
|
self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
|
|
self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
|
|
self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
|
|
self.avgpool = MeanPool2d((7,7))
|
|
self.fc = Dense(512 * block.expansion, num_classes)
|
|
|
|
for m in self.modules():
|
|
if isinstance(m, Conv2d):
|
|
n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
|
|
m.weight.data.normal_(0, math.sqrt(2. / n))
|
|
elif isinstance(m, BatchNorm):
|
|
m.weight.data.fill_(1)
|
|
m.bias.data.zero_()
|
|
|
|
def _make_layer(self, block, planes, blocks, stride=1):
|
|
downsample = None
|
|
if stride != 1 or self.inplanes != planes * block.expansion:
|
|
downsample = SequentialLayer([
|
|
Conv2d(planes * block.expansion,(1,1),(stride,stride),in_channels=self.inplanes),
|
|
BatchNorm(planes * block.expansion)])
|
|
|
|
layers = []
|
|
layers.append(block(self.inplanes, planes, stride, downsample))
|
|
self.inplanes = planes * block.expansion
|
|
for i in range(1, blocks):
|
|
layers.append(block(self.inplanes, planes))
|
|
|
|
return SequentialLayer(layers)
|
|
|
|
def forward(self, x):
|
|
x = self.conv1(x)
|
|
x = self.bn1(x,act='relu')
|
|
x = self.maxpool(x)
|
|
|
|
x = self.layer1(x)
|
|
x = self.layer2(x)
|
|
x = self.layer3(x)
|
|
x = self.layer4(x)
|
|
|
|
x = self.avgpool(x)
|
|
x = x.view(x.size(0), -1)
|
|
x = self.fc(x)
|
|
|
|
return x
|
|
|
|
|
|
|
|
def resnet101():
|
|
"""Constructs a ResNet-101 model.
|
|
Args:
|
|
pretrained (bool): If True, returns a model pre-trained on ImageNet
|
|
"""
|
|
model = ResNet(Bottleneck, [3, 4, 23, 3])
|
|
return model
|
|
|
|
|
|
class resnet(Network):
|
|
def __init__(self, classes, num_layers=101, pretrained=False, class_agnostic=False):
|
|
self.model_path = '.resnet101-5d3b4d8f.pth'
|
|
self.dout_base_model = 1024
|
|
self.pretrained = pretrained
|
|
self.class_agnostic = class_agnostic
|
|
self.num_layers = num_layers
|
|
|
|
Network.__init__(self, classes, class_agnostic)
|
|
|
|
def _init_modules(self):
|
|
resnet = resnet101()
|
|
|
|
if self.pretrained == True:
|
|
print("Loading pretrained weights from %s" %(self.model_path))
|
|
state_dict = torch.load(self.model_path)
|
|
resnet.set_state_dict({k:v for k,v in state_dict.items() if k in resnet.state_dict()})
|
|
|
|
|
|
# Build resnet.
|
|
self.RCNN_base = SequentialLayer([resnet.conv1, resnet.bn1,resnet.relu,
|
|
resnet.maxpool,resnet.layer1,resnet.layer2,resnet.layer3])
|
|
|
|
self.RCNN_top = SequentialLayer([resnet.layer4])
|
|
|
|
self.RCNN_cls_score = Dense(2048, self.n_classes)
|
|
if self.class_agnostic:
|
|
self.RCNN_bbox_pred = Dense(2048, 4)
|
|
else:
|
|
self.RCNN_bbox_pred = Dense(2048, 4 * self.n_classes)
|
|
|
|
# Fix blocks
|
|
for p in self.RCNN_base[0].parameters(): p.requires_grad=False
|
|
for p in self.RCNN_base[1].parameters(): p.requires_grad=False
|
|
|
|
assert (0 <= cfg.RESNET.FIXED_BLOCKS < 4)
|
|
if cfg.RESNET.FIXED_BLOCKS >= 3:
|
|
for p in self.RCNN_base[6].parameters(): p.requires_grad=False
|
|
if cfg.RESNET.FIXED_BLOCKS >= 2:
|
|
for p in self.RCNN_base[5].parameters(): p.requires_grad=False
|
|
if cfg.RESNET.FIXED_BLOCKS >= 1:
|
|
for p in self.RCNN_base[4].parameters(): p.requires_grad=False
|
|
|
|
def set_bn_fix(m):
|
|
classname = m.__class__.__name__
|
|
if classname.find('BatchNorm') != -1:
|
|
for p in m.parameters(): p.requires_grad=False
|
|
|
|
self.RCNN_base.apply(set_bn_fix)
|
|
self.RCNN_top.apply(set_bn_fix)
|
|
|
|
def train(self, mode=True):
|
|
# Override train so that the training mode is set as we want
|
|
Module.train(self, mode)
|
|
if mode:
|
|
# Set fixed blocks to be in eval mode
|
|
self.RCNN_base.eval()
|
|
self.RCNN_base[5].train()
|
|
self.RCNN_base[6].train()
|
|
|
|
def set_bn_eval(m):
|
|
classname = m.__class__.__name__
|
|
if classname.find('BatchNorm') != -1:
|
|
m.eval()
|
|
|
|
self.RCNN_base.apply(set_bn_eval)
|
|
self.RCNN_top.apply(set_bn_eval)
|
|
|
|
def _head_to_tail(self, pool5):
|
|
fc7 = self.RCNN_top(pool5).mean(3).mean(2)
|
|
return fc7
|