def __init__(self, in_planes, planes, stride=1, group_norm=0):
super(PreActBottleneck, self).__init__()
self.bn1 = norm2d(in_planes, group_norm)
self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=1, bias=False)
self.bn2 = norm2d(planes, group_norm)
self.gate1 = GateLayer(planes, planes, [1, -1, 1, 1])
self.conv2 = nn.Conv2d(planes,
planes,
kernel_size=3,
stride=stride,
padding=1,
bias=False)
self.bn3 = norm2d(planes, group_norm)
self.gate2 = GateLayer(planes, planes, [1, -1, 1, 1])
self.conv3 = nn.Conv2d(planes,
self.expansion * planes,
kernel_size=1,
bias=False)
self.gate3 = GateLayer(self.expansion * planes,
self.expansion * planes, [1, -1, 1, 1])
if stride != 1 or in_planes != self.expansion * planes:
self.shortcut = nn.Sequential(
nn.Conv2d(in_planes,
self.expansion * planes,
kernel_size=1,
stride=stride,
bias=False))
self.gate_shortcut = GateLayer(self.expansion * planes,
self.expansion * planes,
[1, -1, 1, 1])
def __init__(self, inplanes, planes, stride=1, downsample=None, gate=None):
super(BasicBlock, self).__init__()
self.conv1 = conv3x3(inplanes, planes, stride)
self.bn1 = nn.BatchNorm2d(planes)
self.relu = nn.ReLU(inplace=True)
self.gate1 = GateLayer(planes, planes, [1, -1, 1, 1])
self.conv2 = conv3x3(planes, planes)
self.bn2 = nn.BatchNorm2d(planes)
self.gate2 = GateLayer(planes, planes, [1, -1, 1, 1])
self.downsample = downsample
self.stride = stride
self.gate = gate
def flatten_model(old_net):
"""Removes nested modules. Only works for VGG."""
from collections import OrderedDict
module_list, counter, inserted_view = [], 0, False
gate_counter = 0
print("printing network")
print(" Hard codded network in vgg_bn.py")
for m_indx, module in enumerate(old_net.modules()):
if not isinstance(module, (nn.Sequential, VGG)):
print(m_indx, module)
if isinstance(module, nn.Linear) and not inserted_view:
module_list.append(('flatten', LinView()))
inserted_view = True
# features.0
# classifier
prefix = "features"
if m_indx > 30:
prefix = "classifier"
if m_indx == 32:
counter = 0
# prefix = ""
module_list.append((prefix + str(counter), module))
if isinstance(module, nn.BatchNorm2d):
planes = module.num_features
gate = GateLayer(planes, planes, [1, -1, 1, 1])
module_list.append(('gate%d' % (gate_counter), gate))
print("gate ", counter, planes)
gate_counter += 1
if isinstance(module, nn.BatchNorm1d):
planes = module.num_features
gate = GateLayer(planes, planes, [1, -1])
module_list.append(('gate%d' % (gate_counter), gate))
print("gate ", counter, planes)
gate_counter += 1
counter += 1
new_net = nn.Sequential(OrderedDict(module_list))
return new_net
def __init__(self, dataset="CIFAR10"):
super(LeNet, self).__init__()
if dataset == "CIFAR10":
nunits_input = 3
nuintis_fc = 32 * 5 * 5
elif dataset == "MNIST":
nunits_input = 1
nuintis_fc = 32 * 4 * 4
self.conv1 = nn.Conv2d(nunits_input, 16, 5)
self.gate1 = GateLayer(16, 16, [1, -1, 1, 1])
self.conv2 = nn.Conv2d(16, 32, 5)
self.gate2 = GateLayer(32, 32, [1, -1, 1, 1])
self.fc1 = nn.Linear(nuintis_fc, 120)
self.gate3 = GateLayer(120, 120, [1, -1])
self.fc2 = nn.Linear(120, 84)
self.gate4 = GateLayer(84, 84, [1, -1])
self.fc3 = nn.Linear(84, 10)
def __init__(self, inplanes, planes, stride=1, downsample=None, gate=None):
super(Bottleneck, self).__init__()
self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
self.bn1 = nn.BatchNorm2d(planes)
self.gate1 = GateLayer(planes,planes,[1, -1, 1, 1])
self.relu1 = nn.ReLU(inplace=True)
self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
padding=1, bias=False)
self.bn2 = nn.BatchNorm2d(planes)
self.gate2 = GateLayer(planes,planes,[1, -1, 1, 1])
self.relu2 = nn.ReLU(inplace=True)
self.conv3 = nn.Conv2d(planes, planes * self.expansion, kernel_size=1, bias=False)
self.bn3 = nn.BatchNorm2d(planes * self.expansion)
self.relu3 = nn.ReLU(inplace=True)
self.downsample = downsample
self.stride = stride
self.gate = gate
def __init__(self, num_input_features, num_output_features, gate_types):
super(_Transition, self).__init__()
self.add_module('norm', nn.BatchNorm2d(num_input_features))
self.add_module('relu', nn.ReLU(inplace=True))
if 'input' in gate_types:
self.add_module(
'gate): (input',
GateLayer(num_input_features, num_input_features,
[1, -1, 1, 1]))
self.add_module(
'conv',
nn.Conv2d(num_input_features,
num_output_features,
kernel_size=1,
stride=1,
bias=False))
self.add_module('pool', nn.AvgPool2d(kernel_size=2, stride=2))
if 'output_conv' in gate_types:
self.add_module(
'gate): (output_conv',
GateLayer(num_output_features, num_output_features,
[1, -1, 1, 1]))
def __init__(self, num_input_features, growth_rate, bn_size, drop_rate,
gate_types):
super(_DenseLayer, self).__init__()
self.add_module('norm1', nn.BatchNorm2d(num_input_features)),
self.add_module('relu1', nn.ReLU(inplace=True)),
if 'input' in gate_types:
self.add_module(
'gate1): (input',
GateLayer(num_input_features, num_input_features,
[1, -1, 1, 1]))
self.add_module(
'conv1',
nn.Conv2d(num_input_features,
bn_size * growth_rate,
kernel_size=1,
stride=1,
bias=False)),
self.add_module('norm2', nn.BatchNorm2d(bn_size * growth_rate)),
if 'output_bn' in gate_types:
self.add_module(
'gate2): (output_bn',
GateLayer(bn_size * growth_rate, bn_size * growth_rate,
[1, -1, 1, 1]))
self.add_module('relu2', nn.ReLU(inplace=True)),
self.add_module(
'conv2',
nn.Conv2d(bn_size * growth_rate,
growth_rate,
kernel_size=3,
stride=1,
padding=1,
bias=False)),
if 'output_conv' in gate_types:
self.add_module('gate3): (output_conv',
GateLayer(growth_rate, growth_rate, [1, -1, 1, 1]))
self.drop_rate = drop_rate
def __init__(self, block, layers, num_classes=1000, skip_gate = True):
self.inplanes = 64
super(ResNet, self).__init__()
self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3,
bias=False)
self.bn1 = nn.BatchNorm2d(64)
self.relu = nn.ReLU(inplace=True)
self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
gate = skip_gate
self.gate = gate
if gate:
# self.gate_skip1 = GateLayer(64,64,[1, -1, 1, 1])
self.gate_skip64 = GateLayer(64*4,64*4,[1, -1, 1, 1])
self.gate_skip128 = GateLayer(128*4,128*4,[1, -1, 1, 1])
self.gate_skip256 = GateLayer(256*4,256*4,[1, -1, 1, 1])
self.gate_skip512 = GateLayer(512*4,512*4,[1, -1, 1, 1])
if block == BasicBlock:
self.gate_skip64 = GateLayer(64, 64, [1, -1, 1, 1])
self.gate_skip128 = GateLayer(128, 128, [1, -1, 1, 1])
self.gate_skip256 = GateLayer(256, 256, [1, -1, 1, 1])
self.gate_skip512 = GateLayer(512, 512, [1, -1, 1, 1])
else:
self.gate_skip64 = None
self.gate_skip128 = None
self.gate_skip256 = None
self.gate_skip512 = None
self.layer1 = self._make_layer(block, 64, layers[0], gate = self.gate_skip64)
self.layer2 = self._make_layer(block, 128, layers[1], stride=2, gate=self.gate_skip128)
self.layer3 = self._make_layer(block, 256, layers[2], stride=2, gate=self.gate_skip256)
self.layer4 = self._make_layer(block, 512, layers[3], stride=2, gate=self.gate_skip512)
self.avgpool = nn.AvgPool2d(7, stride=1)
self.fc = nn.Linear(512 * block.expansion, num_classes)
for m in self.modules():
if isinstance(m, nn.Conv2d):
nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
elif isinstance(m, nn.BatchNorm2d):
nn.init.constant_(m.weight, 1)
nn.init.constant_(m.bias, 0)
def __init__(
self,
growth_rate=32,
block_config=(6, 12, 24, 16),
num_init_features=64,
bn_size=4,
drop_rate=0,
num_classes=1000,
gate_types=['input', 'output_bn', 'output_conv', 'bottom', 'top']):
super(DenseNet, self).__init__()
# First convolution
self.features = nn.Sequential(
OrderedDict([
('conv0',
nn.Conv2d(3,
num_init_features,
kernel_size=7,
stride=2,
padding=3,
bias=False)),
('norm0', nn.BatchNorm2d(num_init_features)),
('relu0', nn.ReLU(inplace=True)),
('pool0', nn.MaxPool2d(kernel_size=3, stride=2, padding=1)),
]))
if 'bottom' in gate_types:
self.features.add_module(
'gate0): (bottom',
GateLayer(num_init_features, num_init_features, [1, -1, 1, 1]))
# Each denseblock
num_features = num_init_features
for i, num_layers in enumerate(block_config):
block = _DenseBlock(num_layers=num_layers,
num_input_features=num_features,
bn_size=bn_size,
growth_rate=growth_rate,
drop_rate=drop_rate,
gate_types=gate_types)
self.features.add_module('denseblock%d' % (i + 1), block)
num_features = num_features + num_layers * growth_rate
if i != len(block_config) - 1:
trans = _Transition(num_input_features=num_features,
num_output_features=num_features // 2,
gate_types=gate_types)
self.features.add_module('transition%d' % (i + 1), trans)
num_features = num_features // 2
# Final batch norm
self.features.add_module('norm5', nn.BatchNorm2d(num_features))
if 'top' in gate_types:
self.features.add_module(
'gate5): (top',
GateLayer(num_features, num_features, [1, -1, 1, 1]))
# Linear layer
self.classifier = nn.Linear(num_features, num_classes)
# Official init from torch repo.
for m in self.modules():
if isinstance(m, nn.Conv2d):
nn.init.kaiming_normal_(m.weight)
elif isinstance(m, nn.BatchNorm2d):
nn.init.constant_(m.weight, 1)
nn.init.constant_(m.bias, 0)
elif isinstance(m, nn.Linear):
nn.init.constant_(m.bias, 0)
def __init__(self,
block,
num_blocks,
num_classes=10,
group_norm=0,
dataset="CIFAR10"):
super(PreActResNet, self).__init__()
self.in_planes = 64
self.dataset = dataset
if dataset == "CIFAR10":
self.conv1 = nn.Conv2d(3,
64,
kernel_size=3,
stride=1,
padding=1,
bias=False)
num_classes = 10
elif dataset == "Imagenet":
self.conv1 = nn.Conv2d(3,
64,
kernel_size=7,
stride=2,
padding=3,
bias=False)
self.bn1 = nn.BatchNorm2d(64)
self.relu = nn.ReLU(inplace=True)
self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
num_classes = 1000
self.gate_in = GateLayer(64, 64, [1, -1, 1, 1])
self.layer1 = self._make_layer(block,
64,
num_blocks[0],
stride=1,
group_norm=group_norm)
self.layer2 = self._make_layer(block,
128,
num_blocks[1],
stride=2,
group_norm=group_norm)
self.layer3 = self._make_layer(block,
256,
num_blocks[2],
stride=2,
group_norm=group_norm)
self.layer4 = self._make_layer(block,
512,
num_blocks[3],
stride=2,
group_norm=group_norm)
if dataset == "CIFAR10":
self.avgpool = nn.AvgPool2d(4, stride=1)
self.linear = nn.Linear(512 * block.expansion, num_classes)
elif dataset == "Imagenet":
self.avgpool = nn.AvgPool2d(7, stride=1)
self.fc = nn.Linear(512 * block.expansion, num_classes)
for m in self.modules():
if isinstance(m, nn.Conv2d):
nn.init.kaiming_normal_(m.weight,
mode='fan_out',
nonlinearity='relu')
if isinstance(m, nn.BatchNorm2d):
m.bias.data.zero_()
