def __init__(self, capacity, quant_schemes):
super(MLP, self).__init__()
nh = int(2048 * capacity)
self.phi1_fc = StochasticLinear(*quant_schemes['phi1_fc'], 28 * 28, nh, bias=False)
self.phi1_bn = nn.BatchNorm1d(nh)
self.phi1_act = StochasticActivation(*quant_schemes['phi1_act'])
self.phi2_fc = StochasticLinear(*quant_schemes['phi2_fc'], nh, nh, bias=False)
self.phi2_bn = nn.BatchNorm1d(nh)
self.phi2_act = StochasticActivation(*quant_schemes['phi2_act'])
self.phi3_fc = StochasticLinear(*quant_schemes['phi3_fc'], nh, nh, bias=False)
self.phi3_bn = nn.BatchNorm1d(nh)
self.phi3_act = StochasticActivation(*quant_schemes['phi3_act'])
self.phi4_fc = StochasticLinear(*quant_schemes['phi4_fc'], nh, 10, bias=False)
self.phi4_bn = nn.BatchNorm1d(10)
def activ(name, nc):
if quantAct:
if quantActSTENumLevels != None and quantActSTENumLevels > 0:
return STEActivation(startEpoch=0,
numLevels=quantActSTENumLevels)
else:
return StochasticActivation(*quant_schemes[name], nc)
else:
assert(quantActSTENumLevels == None or quantActSTENumLevels <= 0)
return nn.ReLU(inplace=True)
def __init__(self, capacity, quant_schemes):
super(AlexNet, self).__init__()
c0 = 3
c1 = int(64 * capacity)
c2 = int(64 * 3 * capacity)
c3 = int(64 * 6 * capacity)
c4 = int(64 * 4 * capacity)
c5 = 256
nh = 4096
# convolutional layers
self.phi1_conv = StochasticConv2d(*quant_schemes['phi1_conv'], c0, c1, kernel_size=11, stride=4, padding=2, bias=False)
self.phi1_mp = nn.MaxPool2d(kernel_size=3, stride=2)
self.phi1_bn = nn.BatchNorm2d(c1)
self.phi1_act = StochasticActivation(*quant_schemes['phi1_act'])
self.phi2_conv = StochasticConv2d(*quant_schemes['phi2_conv'], c1, c2, kernel_size=5, padding=2, bias=False)
self.phi2_mp = nn.MaxPool2d(kernel_size=3, stride=2)
self.phi2_bn = nn.BatchNorm2d(c2)
self.phi2_act = StochasticActivation(*quant_schemes['phi2_act'])
self.phi3_conv = StochasticConv2d(*quant_schemes['phi3_conv'], c2, c3, kernel_size=3, padding=1, bias=False)
self.phi3_bn = nn.BatchNorm2d(c3)
self.phi3_act = StochasticActivation(*quant_schemes['phi3_act'])
self.phi4_conv = StochasticConv2d(*quant_schemes['phi4_conv'], c3, c4, kernel_size=3, padding=1, bias=False)
self.phi4_bn = nn.BatchNorm2d(c4)
self.phi4_act = StochasticActivation(*quant_schemes['phi4_act'])
self.phi5_conv = StochasticConv2d(*quant_schemes['phi5_conv'], c4, c5, kernel_size=3, padding=1, bias=False)
self.phi5_mp = nn.MaxPool2d(kernel_size=3, stride=2)
self.phi5_bn = nn.BatchNorm2d(c5)
self.phi5_act = StochasticActivation(*quant_schemes['phi5_act'])
# fully connected layers
self.phi6_fc = StochasticLinear(*quant_schemes['phi6_fc'], c5 * 6 * 6, nh, bias=False)
self.phi6_bn = nn.BatchNorm1d(nh)
self.phi6_act = StochasticActivation(*quant_schemes['phi6_act'])
self.phi7_fc = StochasticLinear(*quant_schemes['phi7_fc'], nh, nh, bias=False)
self.phi7_bn = nn.BatchNorm1d(nh)
self.phi7_act = StochasticActivation(*quant_schemes['phi7_act'])
self.phi8_fc = StochasticLinear(*quant_schemes['phi8_fc'], nh, 1000, bias=False)
self.phi8_bn = nn.BatchNorm1d(1000)
def __init__(self, capacity, expansion, quant_schemes):
super(MobileNetv2Residuals, self).__init__()
c0 = 3
t0 = int(32 * capacity) * 1
c1 = int(16 * capacity)
t1 = c1 * expansion
c2 = int(24 * capacity)
t2 = c2 * expansion
c3 = int(32 * capacity)
t3 = c3 * expansion
c4 = int(64 * capacity)
t4 = c4 * expansion
c5 = int(96 * capacity)
t5 = c5 * expansion
c6 = int(160 * capacity)
t6 = c6 * expansion
c7 = int(320 * capacity)
c8 = max(int(1280 * capacity), 1280)
# first block
self.phi01_conv = nn.Conv2d(c0,
t0,
kernel_size=3,
stride=2,
padding=1,
bias=False)
self.phi01_bn = nn.BatchNorm2d(t0)
self.phi01_act = nn.ReLU6(inplace=True)
self.phi02_conv = nn.Conv2d(t0,
t0,
kernel_size=3,
stride=1,
padding=1,
groups=t0,
bias=False)
self.phi02_bn = nn.BatchNorm2d(t0)
self.phi02_act = nn.ReLU6(inplace=True)
self.phi03_conv = nn.Conv2d(t0,
c1,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi03_bn = nn.BatchNorm2d(c1)
# second block
self.phi04_conv = nn.Conv2d(c1,
t1,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi04_bn = nn.BatchNorm2d(t1)
self.phi04_act = nn.ReLU6(inplace=True)
self.phi05_conv = nn.Conv2d(t1,
t1,
kernel_size=3,
stride=2,
padding=1,
groups=t1,
bias=False)
self.phi05_bn = nn.BatchNorm2d(t1)
self.phi05_act = nn.ReLU6(inplace=True)
self.phi06_conv = nn.Conv2d(t1,
c2,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi06_bn = nn.BatchNorm2d(c2)
self.phi06_act = StochasticActivation(*quant_schemes['phi06_act'])
self.phi07_conv = StochasticConv2d(*quant_schemes['phi07_conv'],
c2,
t2,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi07_bn = nn.BatchNorm2d(t2)
self.phi07_act = StochasticActivation(*quant_schemes['phi07_act'])
self.phi08_conv = StochasticConv2d(*quant_schemes['phi08_conv'],
t2,
t2,
kernel_size=3,
stride=1,
padding=1,
groups=t2,
bias=False)
self.phi08_bn = nn.BatchNorm2d(t2)
self.phi08_act = StochasticActivation(*quant_schemes['phi08_act'])
self.phi09_conv = StochasticConv2d(*quant_schemes['phi09_conv'],
t2,
c2,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi09_bn = nn.BatchNorm2d(c2)
# third block
self.phi10_conv = nn.Conv2d(c2,
t2,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi10_bn = nn.BatchNorm2d(t2)
self.phi10_act = nn.ReLU6(inplace=True)
self.phi11_conv = nn.Conv2d(t2,
t2,
kernel_size=3,
stride=2,
padding=1,
groups=t2,
bias=False)
self.phi11_bn = nn.BatchNorm2d(t2)
self.phi11_act = nn.ReLU6(inplace=True)
self.phi12_conv = nn.Conv2d(t2,
c3,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi12_bn = nn.BatchNorm2d(c3)
self.phi12_act = StochasticActivation(*quant_schemes['phi12_act'])
self.phi13_conv = StochasticConv2d(*quant_schemes['phi13_conv'],
c3,
t3,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi13_bn = nn.BatchNorm2d(t3)
self.phi13_act = StochasticActivation(*quant_schemes['phi13_act'])
self.phi14_conv = StochasticConv2d(*quant_schemes['phi14_conv'],
t3,
t3,
kernel_size=3,
stride=1,
padding=1,
groups=t3,
bias=False)
self.phi14_bn = nn.BatchNorm2d(t3)
self.phi14_act = StochasticActivation(*quant_schemes['phi14_act'])
self.phi15_conv = StochasticConv2d(*quant_schemes['phi15_conv'],
t3,
c3,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi15_bn = nn.BatchNorm2d(c3)
self.phi15_act = StochasticActivation(*quant_schemes['phi15_act'])
self.phi16_conv = StochasticConv2d(*quant_schemes['phi16_conv'],
c3,
t3,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi16_bn = nn.BatchNorm2d(t3)
self.phi16_act = StochasticActivation(*quant_schemes['phi16_act'])
self.phi17_conv = StochasticConv2d(*quant_schemes['phi17_conv'],
t3,
t3,
kernel_size=3,
stride=1,
padding=1,
groups=t3,
bias=False)
self.phi17_bn = nn.BatchNorm2d(t3)
self.phi17_act = StochasticActivation(*quant_schemes['phi17_act'])
self.phi18_conv = StochasticConv2d(*quant_schemes['phi18_conv'],
t3,
c3,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi18_bn = nn.BatchNorm2d(c3)
# fourth block
self.phi19_conv = nn.Conv2d(c3,
t3,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi19_bn = nn.BatchNorm2d(t3)
self.phi19_act = nn.ReLU6(inplace=True)
self.phi20_conv = nn.Conv2d(t3,
t3,
kernel_size=3,
stride=2,
padding=1,
groups=t3,
bias=False)
self.phi20_bn = nn.BatchNorm2d(t3)
self.phi20_act = nn.ReLU6(inplace=True)
self.phi21_conv = nn.Conv2d(t3,
c4,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi21_bn = nn.BatchNorm2d(c4)
self.phi21_act = StochasticActivation(*quant_schemes['phi21_act'])
self.phi22_conv = StochasticConv2d(*quant_schemes['phi22_conv'],
c4,
t4,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi22_bn = nn.BatchNorm2d(t4)
self.phi22_act = StochasticActivation(*quant_schemes['phi22_act'])
self.phi23_conv = StochasticConv2d(*quant_schemes['phi23_conv'],
t4,
t4,
kernel_size=3,
stride=1,
padding=1,
groups=t4,
bias=False)
self.phi23_bn = nn.BatchNorm2d(t4)
self.phi23_act = StochasticActivation(*quant_schemes['phi23_act'])
self.phi24_conv = StochasticConv2d(*quant_schemes['phi24_conv'],
t4,
c4,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi24_bn = nn.BatchNorm2d(c4)
self.phi24_act = StochasticActivation(*quant_schemes['phi24_act'])
self.phi25_conv = StochasticConv2d(*quant_schemes['phi25_conv'],
c4,
t4,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi25_bn = nn.BatchNorm2d(t4)
self.phi25_act = StochasticActivation(*quant_schemes['phi25_act'])
self.phi26_conv = StochasticConv2d(*quant_schemes['phi26_conv'],
t4,
t4,
kernel_size=3,
stride=1,
padding=1,
groups=t4,
bias=False)
self.phi26_bn = nn.BatchNorm2d(t4)
self.phi26_act = StochasticActivation(*quant_schemes['phi26_act'])
self.phi27_conv = StochasticConv2d(*quant_schemes['phi27_conv'],
t4,
c4,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi27_bn = nn.BatchNorm2d(c4)
self.phi27_act = StochasticActivation(*quant_schemes['phi27_act'])
self.phi28_conv = StochasticConv2d(*quant_schemes['phi28_conv'],
c4,
t4,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi28_bn = nn.BatchNorm2d(t4)
self.phi28_act = StochasticActivation(*quant_schemes['phi28_act'])
self.phi29_conv = StochasticConv2d(*quant_schemes['phi29_conv'],
t4,
t4,
kernel_size=3,
stride=1,
padding=1,
groups=t4,
bias=False)
self.phi29_bn = nn.BatchNorm2d(t4)
self.phi29_act = StochasticActivation(*quant_schemes['phi29_act'])
self.phi30_conv = StochasticConv2d(*quant_schemes['phi30_conv'],
t4,
c4,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi30_bn = nn.BatchNorm2d(c4)
# fifth block
self.phi31_conv = nn.Conv2d(c4,
t4,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi31_bn = nn.BatchNorm2d(t4)
self.phi31_act = nn.ReLU6(inplace=True)
self.phi32_conv = nn.Conv2d(t4,
t4,
kernel_size=3,
stride=1,
padding=1,
groups=t4,
bias=False)
self.phi32_bn = nn.BatchNorm2d(t4)
self.phi32_act = nn.ReLU6(inplace=True)
self.phi33_conv = nn.Conv2d(t4,
c5,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi33_bn = nn.BatchNorm2d(c5)
self.phi33_act = StochasticActivation(*quant_schemes['phi33_act'])
self.phi34_conv = StochasticConv2d(*quant_schemes['phi34_conv'],
c5,
t5,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi34_bn = nn.BatchNorm2d(t5)
self.phi34_act = StochasticActivation(*quant_schemes['phi34_act'])
self.phi35_conv = StochasticConv2d(*quant_schemes['phi35_conv'],
t5,
t5,
kernel_size=3,
stride=1,
padding=1,
groups=t5,
bias=False)
self.phi35_bn = nn.BatchNorm2d(t5)
self.phi35_act = StochasticActivation(*quant_schemes['phi35_act'])
self.phi36_conv = StochasticConv2d(*quant_schemes['phi36_conv'],
t5,
c5,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi36_bn = nn.BatchNorm2d(c5)
self.phi36_act = StochasticActivation(*quant_schemes['phi36_act'])
self.phi37_conv = StochasticConv2d(*quant_schemes['phi37_conv'],
c5,
t5,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi37_bn = nn.BatchNorm2d(t5)
self.phi37_act = StochasticActivation(*quant_schemes['phi37_act'])
self.phi38_conv = StochasticConv2d(*quant_schemes['phi38_conv'],
t5,
t5,
kernel_size=3,
stride=1,
padding=1,
groups=t5,
bias=False)
self.phi38_bn = nn.BatchNorm2d(t5)
self.phi38_act = StochasticActivation(*quant_schemes['phi38_act'])
self.phi39_conv = StochasticConv2d(*quant_schemes['phi39_conv'],
t5,
c5,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi39_bn = nn.BatchNorm2d(c5)
# sixth block
self.phi40_conv = nn.Conv2d(c5,
t5,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi40_bn = nn.BatchNorm2d(t5)
self.phi40_act = nn.ReLU6(inplace=True)
self.phi41_conv = nn.Conv2d(t5,
t5,
kernel_size=3,
stride=2,
padding=1,
groups=t5,
bias=False)
self.phi41_bn = nn.BatchNorm2d(t5)
self.phi41_act = nn.ReLU6(inplace=True)
self.phi42_conv = nn.Conv2d(t5,
c6,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi42_bn = nn.BatchNorm2d(c6)
self.phi42_act = StochasticActivation(*quant_schemes['phi42_act'])
self.phi43_conv = StochasticConv2d(*quant_schemes['phi43_conv'],
c6,
t6,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi43_bn = nn.BatchNorm2d(t6)
self.phi43_act = StochasticActivation(*quant_schemes['phi43_act'])
self.phi44_conv = StochasticConv2d(*quant_schemes['phi44_conv'],
t6,
t6,
kernel_size=3,
stride=1,
padding=1,
groups=t6,
bias=False)
self.phi44_bn = nn.BatchNorm2d(t6)
self.phi44_act = StochasticActivation(*quant_schemes['phi44_act'])
self.phi45_conv = StochasticConv2d(*quant_schemes['phi45_conv'],
t6,
c6,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi45_bn = nn.BatchNorm2d(c6)
self.phi45_act = StochasticActivation(*quant_schemes['phi45_act'])
self.phi46_conv = StochasticConv2d(*quant_schemes['phi46_conv'],
c6,
t6,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi46_bn = nn.BatchNorm2d(t6)
self.phi46_act = StochasticActivation(*quant_schemes['phi46_act'])
self.phi47_conv = StochasticConv2d(*quant_schemes['phi47_conv'],
t6,
t6,
kernel_size=3,
stride=1,
padding=1,
groups=t6,
bias=False)
self.phi47_bn = nn.BatchNorm2d(t6)
self.phi47_act = StochasticActivation(*quant_schemes['phi47_act'])
self.phi48_conv = StochasticConv2d(*quant_schemes['phi48_conv'],
t6,
c6,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi48_bn = nn.BatchNorm2d(c6)
# seventh block
self.phi49_conv = nn.Conv2d(c6,
t6,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi49_bn = nn.BatchNorm2d(t6)
self.phi49_act = nn.ReLU6(inplace=True)
self.phi50_conv = nn.Conv2d(t6,
t6,
kernel_size=3,
stride=1,
padding=1,
groups=t6,
bias=False)
self.phi50_bn = nn.BatchNorm2d(t6)
self.phi50_act = nn.ReLU6(inplace=True)
self.phi51_conv = nn.Conv2d(t6,
c7,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi51_bn = nn.BatchNorm2d(c7)
# classifier
self.phi52_conv = nn.Conv2d(c7,
c8,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi52_bn = nn.BatchNorm2d(c8)
self.phi52_act = nn.ReLU6(inplace=True)
self.phi53_avg = nn.AvgPool2d(kernel_size=7, stride=1, padding=0)
self.phi53_fc = nn.Linear(c8, 1000)
self._initialize_weights()
def __init__(self,
capacity=1,
expansion=6,
quant_schemes=None,
quantAct=True,
quantActSTENumLevels=None,
quantWeights=True,
weightInqSchedule=None,
weightInqBits=2,
weightInqStrategy="magnitude",
quantSkipFirstLayer=False):
super().__init__(capacity, expansion)
c0 = 3
t0 = int(32 * capacity) * 1
c1 = int(16 * capacity)
t1 = c1 * expansion
c2 = int(24 * capacity)
t2 = c2 * expansion
c3 = int(32 * capacity)
t3 = c3 * expansion
c4 = int(64 * capacity)
t4 = c4 * expansion
c5 = int(96 * capacity)
t5 = c5 * expansion
c6 = int(160 * capacity)
t6 = c6 * expansion
c7 = int(320 * capacity)
c8 = max(int(1280 * capacity), 1280)
def activ(name, nc):
if quantAct:
if quantActSTENumLevels != None and quantActSTENumLevels > 0:
return STEActivation(startEpoch=0,
numLevels=quantActSTENumLevels)
else:
return StochasticActivation(*quant_schemes[name], nc)
else:
assert (quantActSTENumLevels == None
or quantActSTENumLevels <= 0)
return nn.ReLU(inplace=True)
def conv2d(name,
ni,
no,
kernel_size=3,
stride=1,
padding=1,
bias=False):
if quantWeights:
if weightInqSchedule == None:
return StochasticConv2d(*quant_schemes[name],
ni,
no,
kernel_size=kernel_size,
stride=stride,
padding=padding,
bias=bias)
else:
return INQConv2d(ni,
no,
kernel_size=kernel_size,
stride=stride,
padding=padding,
bias=bias,
numBits=weightInqBits,
strategy=weightInqStrategy)
else:
return nn.Conv2d(ni,
no,
kernel_size=kernel_size,
stride=stride,
padding=padding,
bias=bias)
def linear(name, ni, no, bias=False):
if quantWeights:
if weightInqSchedule == None:
return StochasticLinear(*quant_schemes[name],
ni,
no,
bias=bias)
else:
return INQLinear(ni,
no,
bias=bias,
numBits=weightInqBits,
strategy=weightInqStrategy)
else:
return nn.Linear(ni, no, bias=bias)
assert (False) # IMPLEMENTATION INCOMPLETE!!!!
# first block
self.phi01_conv = nn.Conv2d(c0,
t0,
kernel_size=3,
stride=2,
padding=1,
bias=False)
self.phi01_bn = nn.BatchNorm2d(t0)
self.phi01_act = nn.ReLU6(inplace=True)
self.phi02_conv = nn.Conv2d(t0,
t0,
kernel_size=3,
stride=1,
padding=1,
groups=t0,
bias=False)
self.phi02_bn = nn.BatchNorm2d(t0)
self.phi02_act = nn.ReLU6(inplace=True)
self.phi03_conv = nn.Conv2d(t0,
c1,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi03_bn = nn.BatchNorm2d(c1)
# second block
self.phi04_conv = nn.Conv2d(c1,
t1,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi04_bn = nn.BatchNorm2d(t1)
self.phi04_act = nn.ReLU6(inplace=True)
self.phi05_conv = nn.Conv2d(t1,
t1,
kernel_size=3,
stride=2,
padding=1,
groups=t1,
bias=False)
self.phi05_bn = nn.BatchNorm2d(t1)
self.phi05_act = nn.ReLU6(inplace=True)
self.phi06_conv = nn.Conv2d(t1,
c2,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi06_bn = nn.BatchNorm2d(c2)
self.phi06_act = StochasticActivation(*quant_schemes['phi06_act'])
self.phi07_conv = StochasticConv2d(*quant_schemes['phi07_conv'],
c2,
t2,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi07_bn = nn.BatchNorm2d(t2)
self.phi07_act = StochasticActivation(*quant_schemes['phi07_act'])
self.phi08_conv = StochasticConv2d(*quant_schemes['phi08_conv'],
t2,
t2,
kernel_size=3,
stride=1,
padding=1,
groups=t2,
bias=False)
self.phi08_bn = nn.BatchNorm2d(t2)
self.phi08_act = StochasticActivation(*quant_schemes['phi08_act'])
self.phi09_conv = StochasticConv2d(*quant_schemes['phi09_conv'],
t2,
c2,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi09_bn = nn.BatchNorm2d(c2)
# third block
self.phi10_conv = nn.Conv2d(c2,
t2,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi10_bn = nn.BatchNorm2d(t2)
self.phi10_act = nn.ReLU6(inplace=True)
self.phi11_conv = nn.Conv2d(t2,
t2,
kernel_size=3,
stride=2,
padding=1,
groups=t2,
bias=False)
self.phi11_bn = nn.BatchNorm2d(t2)
self.phi11_act = nn.ReLU6(inplace=True)
self.phi12_conv = nn.Conv2d(t2,
c3,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi12_bn = nn.BatchNorm2d(c3)
self.phi12_act = StochasticActivation(*quant_schemes['phi12_act'])
self.phi13_conv = StochasticConv2d(*quant_schemes['phi13_conv'],
c3,
t3,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi13_bn = nn.BatchNorm2d(t3)
self.phi13_act = StochasticActivation(*quant_schemes['phi13_act'])
self.phi14_conv = StochasticConv2d(*quant_schemes['phi14_conv'],
t3,
t3,
kernel_size=3,
stride=1,
padding=1,
groups=t3,
bias=False)
self.phi14_bn = nn.BatchNorm2d(t3)
self.phi14_act = StochasticActivation(*quant_schemes['phi14_act'])
self.phi15_conv = StochasticConv2d(*quant_schemes['phi15_conv'],
t3,
c3,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi15_bn = nn.BatchNorm2d(c3)
self.phi15_act = StochasticActivation(*quant_schemes['phi15_act'])
self.phi16_conv = StochasticConv2d(*quant_schemes['phi16_conv'],
c3,
t3,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi16_bn = nn.BatchNorm2d(t3)
self.phi16_act = StochasticActivation(*quant_schemes['phi16_act'])
self.phi17_conv = StochasticConv2d(*quant_schemes['phi17_conv'],
t3,
t3,
kernel_size=3,
stride=1,
padding=1,
groups=t3,
bias=False)
self.phi17_bn = nn.BatchNorm2d(t3)
self.phi17_act = StochasticActivation(*quant_schemes['phi17_act'])
self.phi18_conv = StochasticConv2d(*quant_schemes['phi18_conv'],
t3,
c3,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi18_bn = nn.BatchNorm2d(c3)
# fourth block
self.phi19_conv = nn.Conv2d(c3,
t3,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi19_bn = nn.BatchNorm2d(t3)
self.phi19_act = nn.ReLU6(inplace=True)
self.phi20_conv = nn.Conv2d(t3,
t3,
kernel_size=3,
stride=2,
padding=1,
groups=t3,
bias=False)
self.phi20_bn = nn.BatchNorm2d(t3)
self.phi20_act = nn.ReLU6(inplace=True)
self.phi21_conv = nn.Conv2d(t3,
c4,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi21_bn = nn.BatchNorm2d(c4)
self.phi21_act = StochasticActivation(*quant_schemes['phi21_act'])
self.phi22_conv = StochasticConv2d(*quant_schemes['phi22_conv'],
c4,
t4,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi22_bn = nn.BatchNorm2d(t4)
self.phi22_act = StochasticActivation(*quant_schemes['phi22_act'])
self.phi23_conv = StochasticConv2d(*quant_schemes['phi23_conv'],
t4,
t4,
kernel_size=3,
stride=1,
padding=1,
groups=t4,
bias=False)
self.phi23_bn = nn.BatchNorm2d(t4)
self.phi23_act = StochasticActivation(*quant_schemes['phi23_act'])
self.phi24_conv = StochasticConv2d(*quant_schemes['phi24_conv'],
t4,
c4,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi24_bn = nn.BatchNorm2d(c4)
self.phi24_act = StochasticActivation(*quant_schemes['phi24_act'])
self.phi25_conv = StochasticConv2d(*quant_schemes['phi25_conv'],
c4,
t4,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi25_bn = nn.BatchNorm2d(t4)
self.phi25_act = StochasticActivation(*quant_schemes['phi25_act'])
self.phi26_conv = StochasticConv2d(*quant_schemes['phi26_conv'],
t4,
t4,
kernel_size=3,
stride=1,
padding=1,
groups=t4,
bias=False)
self.phi26_bn = nn.BatchNorm2d(t4)
self.phi26_act = StochasticActivation(*quant_schemes['phi26_act'])
self.phi27_conv = StochasticConv2d(*quant_schemes['phi27_conv'],
t4,
c4,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi27_bn = nn.BatchNorm2d(c4)
self.phi27_act = StochasticActivation(*quant_schemes['phi27_act'])
self.phi28_conv = StochasticConv2d(*quant_schemes['phi28_conv'],
c4,
t4,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi28_bn = nn.BatchNorm2d(t4)
self.phi28_act = StochasticActivation(*quant_schemes['phi28_act'])
self.phi29_conv = StochasticConv2d(*quant_schemes['phi29_conv'],
t4,
t4,
kernel_size=3,
stride=1,
padding=1,
groups=t4,
bias=False)
self.phi29_bn = nn.BatchNorm2d(t4)
self.phi29_act = StochasticActivation(*quant_schemes['phi29_act'])
self.phi30_conv = StochasticConv2d(*quant_schemes['phi30_conv'],
t4,
c4,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi30_bn = nn.BatchNorm2d(c4)
# fifth block
self.phi31_conv = nn.Conv2d(c4,
t4,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi31_bn = nn.BatchNorm2d(t4)
self.phi31_act = nn.ReLU6(inplace=True)
self.phi32_conv = nn.Conv2d(t4,
t4,
kernel_size=3,
stride=1,
padding=1,
groups=t4,
bias=False)
self.phi32_bn = nn.BatchNorm2d(t4)
self.phi32_act = nn.ReLU6(inplace=True)
self.phi33_conv = nn.Conv2d(t4,
c5,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi33_bn = nn.BatchNorm2d(c5)
self.phi33_act = StochasticActivation(*quant_schemes['phi33_act'])
self.phi34_conv = StochasticConv2d(*quant_schemes['phi34_conv'],
c5,
t5,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi34_bn = nn.BatchNorm2d(t5)
self.phi34_act = StochasticActivation(*quant_schemes['phi34_act'])
self.phi35_conv = StochasticConv2d(*quant_schemes['phi35_conv'],
t5,
t5,
kernel_size=3,
stride=1,
padding=1,
groups=t5,
bias=False)
self.phi35_bn = nn.BatchNorm2d(t5)
self.phi35_act = StochasticActivation(*quant_schemes['phi35_act'])
self.phi36_conv = StochasticConv2d(*quant_schemes['phi36_conv'],
t5,
c5,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi36_bn = nn.BatchNorm2d(c5)
self.phi36_act = StochasticActivation(*quant_schemes['phi36_act'])
self.phi37_conv = StochasticConv2d(*quant_schemes['phi37_conv'],
c5,
t5,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi37_bn = nn.BatchNorm2d(t5)
self.phi37_act = StochasticActivation(*quant_schemes['phi37_act'])
self.phi38_conv = StochasticConv2d(*quant_schemes['phi38_conv'],
t5,
t5,
kernel_size=3,
stride=1,
padding=1,
groups=t5,
bias=False)
self.phi38_bn = nn.BatchNorm2d(t5)
self.phi38_act = StochasticActivation(*quant_schemes['phi38_act'])
self.phi39_conv = StochasticConv2d(*quant_schemes['phi39_conv'],
t5,
c5,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi39_bn = nn.BatchNorm2d(c5)
# sixth block
self.phi40_conv = nn.Conv2d(c5,
t5,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi40_bn = nn.BatchNorm2d(t5)
self.phi40_act = nn.ReLU6(inplace=True)
self.phi41_conv = nn.Conv2d(t5,
t5,
kernel_size=3,
stride=2,
padding=1,
groups=t5,
bias=False)
self.phi41_bn = nn.BatchNorm2d(t5)
self.phi41_act = nn.ReLU6(inplace=True)
self.phi42_conv = nn.Conv2d(t5,
c6,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi42_bn = nn.BatchNorm2d(c6)
self.phi42_act = StochasticActivation(*quant_schemes['phi42_act'])
self.phi43_conv = StochasticConv2d(*quant_schemes['phi43_conv'],
c6,
t6,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi43_bn = nn.BatchNorm2d(t6)
self.phi43_act = StochasticActivation(*quant_schemes['phi43_act'])
self.phi44_conv = StochasticConv2d(*quant_schemes['phi44_conv'],
t6,
t6,
kernel_size=3,
stride=1,
padding=1,
groups=t6,
bias=False)
self.phi44_bn = nn.BatchNorm2d(t6)
self.phi44_act = StochasticActivation(*quant_schemes['phi44_act'])
self.phi45_conv = StochasticConv2d(*quant_schemes['phi45_conv'],
t6,
c6,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi45_bn = nn.BatchNorm2d(c6)
self.phi45_act = StochasticActivation(*quant_schemes['phi45_act'])
self.phi46_conv = StochasticConv2d(*quant_schemes['phi46_conv'],
c6,
t6,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi46_bn = nn.BatchNorm2d(t6)
self.phi46_act = StochasticActivation(*quant_schemes['phi46_act'])
self.phi47_conv = StochasticConv2d(*quant_schemes['phi47_conv'],
t6,
t6,
kernel_size=3,
stride=1,
padding=1,
groups=t6,
bias=False)
self.phi47_bn = nn.BatchNorm2d(t6)
self.phi47_act = StochasticActivation(*quant_schemes['phi47_act'])
self.phi48_conv = StochasticConv2d(*quant_schemes['phi48_conv'],
t6,
c6,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi48_bn = nn.BatchNorm2d(c6)
# seventh block
self.phi49_conv = nn.Conv2d(c6,
t6,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi49_bn = nn.BatchNorm2d(t6)
self.phi49_act = nn.ReLU6(inplace=True)
self.phi50_conv = nn.Conv2d(t6,
t6,
kernel_size=3,
stride=1,
padding=1,
groups=t6,
bias=False)
self.phi50_bn = nn.BatchNorm2d(t6)
self.phi50_act = nn.ReLU6(inplace=True)
self.phi51_conv = nn.Conv2d(t6,
c7,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi51_bn = nn.BatchNorm2d(c7)
# classifier
self.phi52_conv = nn.Conv2d(c7,
c8,
kernel_size=1,
stride=1,
padding=0,
bias=False)
self.phi52_bn = nn.BatchNorm2d(c8)
self.phi52_act = nn.ReLU6(inplace=True)
self.phi53_avg = nn.AvgPool2d(kernel_size=7, stride=1, padding=0)
self.phi53_fc = nn.Linear(c8, 1000)
self._initialize_weights()
def __init__(self, capacityFactor=1.0, version=1,
quantized=True, quant_scheme=None,
quantFirstLast=True, withTwoAct=False, noTimePooling=False):
super().__init__()
self.noTimePooling = noTimePooling
def conv1quant(quant_scheme, ni, no, stride=1, padding=1):
return StochasticConv2d(*quant_scheme, ni, no, kernel_size=1,
stride=stride, padding=0, bias=False)
def conv3quant(quant_scheme, ni, no, stride=1, padding=1):
return StochasticConv2d(*quant_scheme, ni, no, kernel_size=3,
stride=stride, padding=1, bias=False)
def conv1float(quant_scheme, ni, no, stride=1, padding=1):
return nn.Conv2d(ni, no, kernel_size=1,
stride=stride, padding=0, bias=False)
def conv3float(quant_scheme, ni, no, stride=1, padding=1):
return nn.Conv2d(ni, no, kernel_size=3,
stride=stride, padding=1, bias=False)
if quantized:
conv1 = conv1quant
conv3 = conv3quant
activ = lambda quant_scheme, nc: StochasticActivation(*quant_scheme, nc)
if withTwoAct:
activ2 = lambda nc: nn.ReLU(inplace=True)
else:
activ2 = lambda nc: nn.Identity()
quantScheme = lambda s: quant_scheme[s]
else:
conv1 = conv1float
conv3 = conv3float
activ = lambda quant_scheme, nc: nn.ReLU(inplace=True)
activ2 = lambda nc: nn.Identity()
quantScheme = lambda s: None
bnorm = lambda nc: nn.BatchNorm2d(nc)
# bnorm = lambda nc: nn.Identity() # don't forget to enable/disable bias
c = lambda v: math.ceil(v*capacityFactor)
c1, c2, c3, c4, c5, c6 = c(64), c(64), c(128), c(128), c(128), c(128)
if version >= 2:
c1 = c(32)
if quantFirstLast:
self.phi1_conv = conv3(quantScheme('phi1_conv'), 1, c1)
else:
self.phi1_conv = conv3float(None, 1, c1)
self.phi1_act2 = activ2(c1)
self.phi1_bn = bnorm(c1)
self.phi1_act = activ(quantScheme('phi1_act'), c1)
self.phi2_conv = conv3(quantScheme('phi2_conv'), c1, c2, stride=2)
self.phi2_act2 = activ2(c2)
self.phi2_bn = bnorm(c2)
self.phi2_act = activ(quantScheme('phi2_act'), c2)
self.phi3_conv = conv3(quantScheme('phi3_conv'), c2, c3)
self.phi3_act2 = activ2(c3)
self.phi3_bn = bnorm(c3)
self.phi3_act = activ(quantScheme('phi3_act'), c3)
if version >= 3:
self.phi4_do = nn.Dropout2d(0.5)
else:
self.phi4_do = nn.Identity()
self.phi4_conv = conv3(quantScheme('phi4_conv'), c3, c4, stride=2)
self.phi4_act2 = activ2(c4)
self.phi4_bn = bnorm(c4)
self.phi4_act = activ(quantScheme('phi4_act'), c4)
self.phi5_conv = conv3(quantScheme('phi5_conv'), c4, c5)
self.phi5_act2 = activ2(c5)
self.phi5_bn = bnorm(c5)
self.phi5_act = activ(quantScheme('phi5_act'), c5)
self.phi6_conv = conv1(quantScheme('phi6_conv'), c5, c6)
self.phi6_act2 = activ2(c6)
self.phi6_bn = bnorm(c6)
if quantFirstLast:
self.phi6_act = activ(quantScheme('phi6_act'), c6)
self.phi7_conv = conv1(quantScheme('phi7_conv'), c6, 28)
else:
self.phi6_act = nn.Identity()
self.phi7_conv = conv1float(None, c6, 28)
self.phi7_bn = bnorm(28)
if noTimePooling:
self.phi8_pool = nn.AvgPool2d(kernel_size=(16,1), stride=1, padding=0)
else:
self.phi8_pool = nn.AvgPool2d(kernel_size=(16,100), stride=1, padding=0)
def __init__(self, capacity, quant_schemes):
super(VGG, self).__init__()
c0 = 3
c1 = int(128 * capacity)
c2 = int(128 * 2 * capacity)
c3 = int(128 * 4 * capacity)
nh = 1024
# convolutional layers
self.phi1_conv = StochasticConv2d(*quant_schemes['phi1_conv'],
c0,
c1,
kernel_size=3,
stride=1,
padding=1,
bias=False)
self.phi1_bn = nn.BatchNorm2d(c1)
self.phi1_act = StochasticActivation(*quant_schemes['phi1_act'])
self.phi2_conv = StochasticConv2d(*quant_schemes['phi2_conv'],
c1,
c1,
kernel_size=3,
stride=1,
padding=1,
bias=False)
self.phi2_mp = nn.MaxPool2d(kernel_size=2, stride=2)
self.phi2_bn = nn.BatchNorm2d(c1)
self.phi2_act = StochasticActivation(*quant_schemes['phi2_act'])
self.phi3_conv = StochasticConv2d(*quant_schemes['phi3_conv'],
c1,
c2,
kernel_size=3,
stride=1,
padding=1,
bias=False)
self.phi3_bn = nn.BatchNorm2d(c2)
self.phi3_act = StochasticActivation(*quant_schemes['phi3_act'])
self.phi4_conv = StochasticConv2d(*quant_schemes['phi4_conv'],
c2,
c2,
kernel_size=3,
stride=1,
padding=1,
bias=False)
self.phi4_mp = nn.MaxPool2d(kernel_size=2, stride=2)
self.phi4_bn = nn.BatchNorm2d(c2)
self.phi4_act = StochasticActivation(*quant_schemes['phi4_act'])
self.phi5_conv = StochasticConv2d(*quant_schemes['phi5_conv'],
c2,
c3,
kernel_size=3,
stride=1,
padding=1,
bias=False)
self.phi5_bn = nn.BatchNorm2d(c3)
self.phi5_act = StochasticActivation(*quant_schemes['phi5_act'])
self.phi6_conv = StochasticConv2d(*quant_schemes['phi6_conv'],
c3,
c3,
kernel_size=3,
stride=1,
padding=1,
bias=False)
self.phi6_mp = nn.MaxPool2d(kernel_size=2, stride=2)
self.phi6_bn = nn.BatchNorm2d(c3)
self.phi6_act = StochasticActivation(*quant_schemes['phi6_act'])
# fully connected layers
self.phi7_fc = StochasticLinear(*quant_schemes['phi7_fc'],
c3 * 4 * 4,
nh,
bias=False)
self.phi7_bn = nn.BatchNorm1d(nh)
self.phi7_act = StochasticActivation(*quant_schemes['phi7_act'])
self.phi8_fc = StochasticLinear(*quant_schemes['phi8_fc'],
nh,
nh,
bias=False)
self.phi8_bn = nn.BatchNorm1d(nh)
self.phi8_act = StochasticActivation(*quant_schemes['phi8_act'])
self.phi9_fc = StochasticLinear(*quant_schemes['phi9_fc'],
nh,
10,
bias=False)
self.phi9_bn = nn.BatchNorm1d(10)
def activ(name, nc):
if quantAct:
return StochasticActivation(*quant_scheme[name], nc)
else:
return nn.ReLU()