def sparse_conv2d(in_channels,
out_channels,
kernel_size,
stride=1,
padding=0,
dilation=1,
groups=1,
bias=True,
density=1.0):
"""
Get a nn.Conv2d, possibly wrapped in a SparseWeights2d
:param density:
Either a density or a function that returns a density.
:type density: float or function(in_channels, out_channels, kernel_size)
"""
layer = nn.Conv2d(in_channels,
out_channels,
kernel_size,
stride=stride,
padding=padding,
dilation=dilation,
groups=groups,
bias=bias)
if callable(density):
density = density(in_channels, out_channels, kernel_size)
if density < 1.0:
layer = SparseWeights2d(layer, weight_sparsity=density)
return layer
def add_sparse_cnn_layer(
network,
suffix,
in_channels,
out_channels,
use_batch_norm,
weight_sparsity,
percent_on,
k_inference_factor,
boost_strength,
boost_strength_factor,
):
"""Add sparse cnn layer to network.
:param network: The network to add the sparse layer to
:param suffix: Layer suffix. Used to name its components
:param in_channels: input channels
:param out_channels: output channels
:param use_batch_norm: whether or not to use batch norm
:param weight_sparsity: Pct of weights that are allowed to be non-zero
:param percent_on: Pct of ON (non-zero) units
:param k_inference_factor: During inference we increase percent_on by this factor
:param boost_strength: boost strength (0.0 implies no boosting)
:param boost_strength_factor:
boost strength is multiplied by this factor after each epoch
"""
cnn = nn.Conv2d(
in_channels=in_channels,
out_channels=out_channels,
kernel_size=5,
padding=0,
stride=1,
)
if 0 < weight_sparsity < 1.0:
sparse_cnn = SparseWeights2d(cnn, weight_sparsity)
network.add_module("cnnSdr{}_cnn".format(suffix), sparse_cnn)
else:
network.add_module("cnnSdr{}_cnn".format(suffix), cnn)
if use_batch_norm:
bn = nn.BatchNorm2d(out_channels, affine=False)
network.add_module("cnnSdr{}_bn".format(suffix), bn)
# Max pool
maxpool = nn.MaxPool2d(kernel_size=2)
network.add_module("cnnSdr{}_maxpool".format(suffix), maxpool)
if 0 < percent_on < 1.0:
kwinner = KWinners2d(
channels=out_channels,
percent_on=percent_on,
k_inference_factor=k_inference_factor,
boost_strength=boost_strength,
boost_strength_factor=boost_strength_factor,
)
network.add_module("cnnSdr{}_kwinner".format(suffix), kwinner)
else:
network.add_module("cnnSdr{}_relu".format(suffix), nn.ReLU())
def _add_cnn_layer(
self,
index_str,
in_channels,
out_channels,
kernel_size,
percent_on,
weight_sparsity,
k_inference_factor,
boost_strength,
boost_strength_factor,
add_pooling,
use_max_pooling,
):
"""Add a single CNN layer to our modules."""
# Add CNN layer
if kernel_size == 3:
padding = 1
else:
padding = 2
conv2d = nn.Conv2d(in_channels,
out_channels,
kernel_size=kernel_size,
padding=padding)
if weight_sparsity < 1.0:
conv2d = SparseWeights2d(conv2d, weight_sparsity=weight_sparsity)
self.add_module("cnn_" + index_str, conv2d)
self.add_module("bn_" + index_str, nn.BatchNorm2d(out_channels)),
if add_pooling:
if use_max_pooling:
self.add_module("maxpool_" + index_str,
nn.MaxPool2d(kernel_size=2, stride=2))
else:
self.add_module("avgpool_" + index_str,
nn.AvgPool2d(kernel_size=2, stride=2))
if percent_on < 1.0:
self.add_module(
"kwinners_2d_" + index_str,
KWinners2d(
percent_on=percent_on,
channels=out_channels,
k_inference_factor=k_inference_factor,
boost_strength=boost_strength,
boost_strength_factor=boost_strength_factor,
),
)
else:
self.add_module("ReLU_" + index_str, nn.ReLU(inplace=True))
def _sparsify_cnn(parent, cnn_names, weight_sparsity):
"""Enforce weight sparsity on the given cnn modules during training.
:param parent: Parent Layer containing the CNN modules to sparsify
:param cnn_names: List of CNN module names to sparsify
:param weight_sparsity: Percent of weights that are allowed to be non-zero
"""
for i, name in enumerate(cnn_names):
if weight_sparsity[i] >= 1.0:
continue
module = parent.__getattr__(name)
parent.__setattr__(name, SparseWeights2d(module, weight_sparsity[i]))
def test_sparse_weights_2d(self):
in_channels, kernel_size, out_channels = 64, (5, 5), 64
input_size = in_channels * kernel_size[0] * kernel_size[1]
with torch.no_grad():
for percent_on in [0.1, 0.5, 0.9]:
cnn = torch.nn.Conv2d(in_channels=in_channels,
out_channels=out_channels,
kernel_size=kernel_size)
sparse = SparseWeights2d(cnn, percent_on)
nonzeros = torch.nonzero(sparse.module.weight, as_tuple=True)[0]
counts = torch.unique(nonzeros, return_counts=True)[1]
# Expected non-zeros per output channel
expected = [round(input_size * percent_on)] * out_channels
self.assertSequenceEqual(counts.numpy().tolist(), expected)
def test_rezero_after_forward_2d(self):
in_channels, kernel_size, out_channels = 64, (5, 5), 64
input_size = in_channels * kernel_size[0] * kernel_size[1]
with torch.no_grad():
for percent_on in [0.1, 0.5, 0.9]:
cnn = torch.nn.Conv2d(in_channels=in_channels,
out_channels=out_channels,
kernel_size=kernel_size)
sparse = SparseWeights2d(cnn, percent_on)
# Ensure weights are not sparse
sparse.module.weight.data.fill_(1.0)
sparse.train()
x = torch.ones((1,) + (in_channels, kernel_size[0], kernel_size[1]))
sparse(x)
# When training, the forward function should set weights back to zero.
nonzeros = torch.nonzero(sparse.module.weight, as_tuple=True)[0]
counts = torch.unique(nonzeros, return_counts=True)[1]
expected = [round(input_size * percent_on)] * out_channels
self.assertSequenceEqual(counts.numpy().tolist(), expected)
def test_rezero_2d(self):
in_channels, kernel_size, out_channels = 64, (5, 5), 64
input_size = in_channels * kernel_size[0] * kernel_size[1]
with torch.no_grad():
for sparsity in [0.1, 0.5, 0.9]:
cnn = torch.nn.Conv2d(in_channels=in_channels,
out_channels=out_channels,
kernel_size=kernel_size)
sparse = SparseWeights2d(cnn, sparsity=sparsity)
# Ensure weights are not sparse
sparse.module.weight.data.fill_(1.0)
# Rezero, verify the weights become sparse
sparse.rezero_weights()
nonzeros = torch.nonzero(sparse.module.weight,
as_tuple=True)[0]
counts = torch.unique(nonzeros, return_counts=True)[1]
expected = [round(input_size *
(1.0 - sparsity))] * out_channels
self.assertSequenceEqual(counts.numpy().tolist(), expected)
def __init__(
self,
dpc=3,
cnn_w_sparsity=0.05,
linear_w_sparsity=0.5,
cat_w_sparsity=0.01,
n_classes=4,
):
super(ToyNetwork, self).__init__()
conv_channels = 128
self.n_classes = n_classes
self.conv1 = SparseWeights2d(
nn.Conv2d(
in_channels=1,
out_channels=conv_channels,
kernel_size=10,
padding=0,
stride=1,
),
cnn_w_sparsity,
)
self.kwin1 = KWinners2d(conv_channels, percent_on=0.1)
self.bn = nn.BatchNorm2d(conv_channels, affine=False)
self.mp1 = nn.MaxPool2d(kernel_size=2)
self.flatten = Flatten()
self.d1 = DendriteLayer(
in_dim=int(conv_channels / 64) * 7744,
out_dim=1000,
dendrites_per_neuron=dpc,
)
self.linear = SparseWeights(nn.Linear(1000, n_classes + 1),
linear_w_sparsity)
self.cat = SparseWeights(nn.Linear(n_classes + 1, 1000 * dpc),
cat_w_sparsity)
def _add_cnn_layer(self, index_str, in_channels, out_channels, kernel_size,
percent_on, weight_sparsity, add_pooling):
"""
Add a single CNN layer to our modules
"""
# Add CNN layer
if kernel_size == 3:
padding = 1
else:
padding = 2
conv2d = nn.Conv2d(in_channels,
out_channels,
kernel_size=kernel_size,
padding=padding)
if weight_sparsity < 1.0:
conv2d = SparseWeights2d(conv2d, weightSparsity=weight_sparsity)
self.model.add_module("cnn_" + index_str, conv2d)
self.model.add_module("bn_" + index_str, nn.BatchNorm2d(out_channels)),
if add_pooling:
self.model.add_module("avgpool_" + index_str,
nn.AvgPool2d(kernel_size=2))
if percent_on < 1.0:
self.model.add_module(
"kwinners_2d_" + index_str,
KWinners2d(percent_on=percent_on,
channels=out_channels,
kInferenceFactor=self.k_inference_factor,
boostStrength=self.boost_strength,
boostStrengthFactor=self.boost_strength_factor))
else:
self.model.add_module("ReLU_" + index_str, nn.ReLU())
def __init__(self,
cnn_out_channels=(32, 64),
cnn_percent_on=(0.1, 0.2),
cnn_weight_sparsity=(0.6, 0.45),
linear_units=700,
linear_percent_on=0.2,
linear_weight_sparsity=0.2,
boost_strength=1.5,
boost_strength_factor=0.85,
k_inference_factor=1.0,
duty_cycle_period=1000,
kwinner_local=False):
super(MNISTSparseCNN, self).__init__(
OrderedDict([
# First Sparse CNN layer
("cnn1",
SparseWeights2d(nn.Conv2d(1, cnn_out_channels[0], 5),
cnn_weight_sparsity[0])),
("cnn1_maxpool", nn.MaxPool2d(2)),
("cnn1_kwinner",
KWinners2d(channels=cnn_out_channels[0],
percent_on=cnn_percent_on[0],
k_inference_factor=k_inference_factor,
boost_strength=boost_strength,
boost_strength_factor=boost_strength_factor,
duty_cycle_period=duty_cycle_period,
local=kwinner_local)),
# Second Sparse CNN layer
("cnn2",
SparseWeights2d(
nn.Conv2d(cnn_out_channels[0], cnn_out_channels[1], 5),
cnn_weight_sparsity[1])),
("cnn2_maxpool", nn.MaxPool2d(2)),
("cnn2_kwinner",
KWinners2d(channels=cnn_out_channels[1],
percent_on=cnn_percent_on[1],
k_inference_factor=k_inference_factor,
boost_strength=boost_strength,
boost_strength_factor=boost_strength_factor,
duty_cycle_period=duty_cycle_period,
local=kwinner_local)),
("flatten", Flatten()),
# Sparse Linear layer
("linear",
SparseWeights(nn.Linear(16 * cnn_out_channels[1],
linear_units),
weight_sparsity=linear_weight_sparsity)),
("linear_kwinner",
KWinners(n=linear_units,
percent_on=linear_percent_on,
k_inference_factor=k_inference_factor,
boost_strength=boost_strength,
boost_strength_factor=boost_strength_factor,
duty_cycle_period=duty_cycle_period)),
# Classifier
("output", nn.Linear(linear_units, 10)),
("softmax", nn.LogSoftmax(dim=1))
]))
def __init__(self,
cnn_out_channels=(64, 64),
cnn_percent_on=(0.095, 0.125),
cnn_weight_sparsity=(0.5, 0.2),
linear_units=1000,
linear_percent_on=0.1,
linear_weight_sparsity=0.1,
boost_strength=1.5,
boost_strength_factor=0.9,
k_inference_factor=1.0,
duty_cycle_period=1000,
kwinner_local=False):
super(GSCSparseCNN, self).__init__()
# input_shape = (1, 32, 32)
# First Sparse CNN layer
if cnn_weight_sparsity[0] < 1.0:
self.add_module(
"cnn1",
SparseWeights2d(nn.Conv2d(1, cnn_out_channels[0], 5),
weight_sparsity=cnn_weight_sparsity[0]))
else:
self.add_module("cnn1", nn.Conv2d(1, cnn_out_channels[0], 5))
self.add_module("cnn1_batchnorm",
nn.BatchNorm2d(cnn_out_channels[0], affine=False))
self.add_module(
"cnn1_kwinner",
KWinners2d(
channels=cnn_out_channels[0],
percent_on=cnn_percent_on[0],
k_inference_factor=k_inference_factor,
boost_strength=boost_strength,
boost_strength_factor=boost_strength_factor,
duty_cycle_period=duty_cycle_period,
local=kwinner_local,
))
self.add_module("cnn1_maxpool", nn.MaxPool2d(2))
# Second Sparse CNN layer
if cnn_weight_sparsity[1] < 1.0:
self.add_module(
"cnn2",
SparseWeights2d(nn.Conv2d(cnn_out_channels[0],
cnn_out_channels[1], 5),
weight_sparsity=cnn_weight_sparsity[1]))
else:
self.add_module(
"cnn2", nn.Conv2d(cnn_out_channels[0], cnn_out_channels[1], 5))
self.add_module("cnn2_batchnorm",
nn.BatchNorm2d(cnn_out_channels[1], affine=False))
self.add_module(
"cnn2_kwinner",
KWinners2d(
channels=cnn_out_channels[1],
percent_on=cnn_percent_on[1],
k_inference_factor=k_inference_factor,
boost_strength=boost_strength,
boost_strength_factor=boost_strength_factor,
duty_cycle_period=duty_cycle_period,
local=kwinner_local,
))
self.add_module("cnn2_maxpool", nn.MaxPool2d(2))
self.add_module("flatten", Flatten())
# Sparse Linear layer
self.add_module(
"linear",
SparseWeights(nn.Linear(25 * cnn_out_channels[1], linear_units),
weight_sparsity=linear_weight_sparsity))
self.add_module("linear_bn", nn.BatchNorm1d(linear_units,
affine=False))
self.add_module(
"linear_kwinner",
KWinners(n=linear_units,
percent_on=linear_percent_on,
k_inference_factor=k_inference_factor,
boost_strength=boost_strength,
boost_strength_factor=boost_strength_factor,
duty_cycle_period=duty_cycle_period))
# Classifier
self.add_module("output", nn.Linear(linear_units, 12))
self.add_module("softmax", nn.LogSoftmax(dim=1))
def add_sparse_cnn_layer(
network,
suffix,
in_channels,
out_channels,
use_batch_norm,
weight_sparsity,
percent_on,
k_inference_factor,
boost_strength,
boost_strength_factor,
duty_cycle_period,
activation_fct_before_max_pool,
use_kwinners_local,
consolidated_sparse_weights,
):
"""Add sparse cnn layer to network.
:param network: The network to add the sparse layer to
:param suffix: Layer suffix. Used to name its components
:param in_channels: input channels
:param out_channels: output channels
:param use_batch_norm: whether or not to use batch norm
:param weight_sparsity: Pct of weights that are allowed to be non-zero
:param percent_on: Pct of ON (non-zero) units
:param k_inference_factor: During inference we increase percent_on by this factor
:param boost_strength: boost strength (0.0 implies no boosting)
:param boost_strength_factor:
boost strength is multiplied by this factor after each epoch
:param activation_fct_before_max_pool:
If true ReLU/K-winners will be placed before the max_pool step
:param use_kwinners_local:
Whether or not to choose the k-winners 2d locally only across the
channels instead of the whole input
"""
cnn = nn.Conv2d(
in_channels=in_channels,
out_channels=out_channels,
kernel_size=5,
padding=0,
stride=1,
)
if 0 < weight_sparsity < 1.0:
if consolidated_sparse_weights:
sparse_cnn = ConsolidatedSparseWeights2D(cnn, weight_sparsity)
else:
sparse_cnn = SparseWeights2d(cnn, weight_sparsity)
network.add_module("cnn{}_cnn".format(suffix), sparse_cnn)
else:
network.add_module("cnn{}_cnn".format(suffix), cnn)
if use_batch_norm:
bn = nn.BatchNorm2d(out_channels, affine=False)
network.add_module("cnn{}_bn".format(suffix), bn)
if not activation_fct_before_max_pool:
maxpool = nn.MaxPool2d(kernel_size=2)
network.add_module("cnn{}_maxpool".format(suffix), maxpool)
if percent_on >= 1.0 or percent_on <= 0:
network.add_module("cnn{}_relu".format(suffix), nn.ReLU())
else:
kwinner_class = KWinners2dLocal if use_kwinners_local else KWinners2d
kwinner = kwinner_class(
channels=out_channels,
percent_on=percent_on,
k_inference_factor=k_inference_factor,
boost_strength=boost_strength,
boost_strength_factor=boost_strength_factor,
duty_cycle_period=duty_cycle_period,
)
network.add_module("cnn{}_kwinner".format(suffix), kwinner)
if activation_fct_before_max_pool:
maxpool = nn.MaxPool2d(kernel_size=2)
network.add_module("cnn{}_maxpool".format(suffix), maxpool)
def add_sparse_cnn_layer(
network,
suffix,
in_channels,
out_channels,
use_batch_norm,
weight_sparsity,
percent_on,
k_inference_factor,
temperature,
eval_temperature,
temperature_decay_rate,
activation_fct_before_max_pool,
consolidated_sparse_weights,
):
"""Add sparse cnn layer to network.
:param network: The network to add the sparse layer to
:param suffix: Layer suffix. Used to name its components
:param in_channels: input channels
:param out_channels: output channels
:param use_batch_norm: whether or not to use batch norm
:param weight_sparsity: Pct of weights that are allowed to be non-zero
:param percent_on: Pct of ON (non-zero) units
:param k_inference_factor: During inference we increase percent_on by this factor
:param temperature: temperature to use when computing softmax in SampledKWinners.
:param eval_temperature:
temperature to use when computing softmax during evaluation in SampledKWinners.
:param temperature_decay_rate:
amount to decrease the temperature after each epoch of training.
:param activation_fct_before_max_pool:
If true ReLU/K-winners will be placed before the max_pool step
"""
cnn = nn.Conv2d(
in_channels=in_channels,
out_channels=out_channels,
kernel_size=5,
padding=0,
stride=1,
)
if 0 < weight_sparsity < 1.0:
if consolidated_sparse_weights:
sparse_cnn = ConsolidatedSparseWeights2D(cnn, weight_sparsity)
else:
sparse_cnn = SparseWeights2d(cnn, weight_sparsity)
network.add_module("cnn{}_cnn".format(suffix), sparse_cnn)
else:
network.add_module("cnn{}_cnn".format(suffix), cnn)
if use_batch_norm:
bn = nn.BatchNorm2d(out_channels, affine=False)
network.add_module("cnn{}_bn".format(suffix), bn)
if not activation_fct_before_max_pool:
maxpool = nn.MaxPool2d(kernel_size=2)
network.add_module("cnn{}_maxpool".format(suffix), maxpool)
if percent_on >= 1.0 or percent_on <= 0:
network.add_module("cnn{}_relu".format(suffix), nn.ReLU())
else:
kwinner_class = SampledKWinners2d
kwinner = kwinner_class(percent_on=percent_on,
k_inference_factor=k_inference_factor,
temperature=temperature,
eval_temperature=eval_temperature,
temperature_decay_rate=temperature_decay_rate,
relu=False)
network.add_module("cnn{}_kwinner".format(suffix), kwinner)
if activation_fct_before_max_pool:
maxpool = nn.MaxPool2d(kernel_size=2)
network.add_module("cnn{}_maxpool".format(suffix), maxpool)
def __init__(
self,
cnn_out_channels=(64, 64),
cnn_percent_on=(0.095, 0.125),
cnn_weight_sparsity=None,
linear_units=1000,
linear_percent_on=0.1,
linear_weight_sparsity=None,
temperature=10.0,
eval_temperature=1.0,
temperature_decay_rate=0.99,
k_inference_factor=1.0,
cnn_sparsity=(0.5, 0.8),
linear_sparsity=0.9,
):
super(SampledKWinnerGSCSparseCNN, self).__init__()
if cnn_weight_sparsity is not None:
warnings.warn(
"Parameter `cnn_weight_sparsity` is deprecated. Use "
"`cnn_sparsity` instead.",
DeprecationWarning,
)
cnn_sparsity = (1.0 - cnn_weight_sparsity[0], 1.0 - cnn_weight_sparsity[1])
if linear_weight_sparsity is not None:
warnings.warn(
"Parameter `linear_weight_sparsity` is deprecated. Use "
"`linear_sparsity` instead.",
DeprecationWarning,
)
linear_sparsity = 1.0 - linear_weight_sparsity
# input_shape = (1, 32, 32)
# First Sparse CNN layer
if cnn_sparsity[0] > 0:
self.add_module(
"cnn1",
SparseWeights2d(
nn.Conv2d(1, cnn_out_channels[0], 5), sparsity=cnn_sparsity[0]
),
)
else:
self.add_module("cnn1", nn.Conv2d(1, cnn_out_channels[0], 5))
self.add_module(
"cnn1_batchnorm", nn.BatchNorm2d(cnn_out_channels[0], affine=False)
)
self.add_module(
"cnn1_kwinner",
SampledKWinners2d(
percent_on=cnn_percent_on[0],
k_inference_factor=k_inference_factor,
temperature=temperature,
eval_temperature=eval_temperature,
temperature_decay_rate=temperature_decay_rate,
relu=False,
),
)
self.add_module("cnn1_maxpool", nn.MaxPool2d(2))
# Second Sparse CNN layer
if cnn_sparsity[1] > 0:
self.add_module(
"cnn2",
SparseWeights2d(
nn.Conv2d(cnn_out_channels[0], cnn_out_channels[1], 5),
sparsity=cnn_sparsity[1],
),
)
else:
self.add_module(
"cnn2", nn.Conv2d(cnn_out_channels[0], cnn_out_channels[1], 5)
)
self.add_module(
"cnn2_batchnorm", nn.BatchNorm2d(cnn_out_channels[1], affine=False)
)
self.add_module(
"cnn2_kwinner",
SampledKWinners2d(
percent_on=cnn_percent_on[0],
k_inference_factor=k_inference_factor,
temperature=temperature,
eval_temperature=eval_temperature,
temperature_decay_rate=temperature_decay_rate,
relu=False,
),
)
self.add_module("cnn2_maxpool", nn.MaxPool2d(2))
self.add_module("flatten", Flatten())
# Sparse Linear layer
self.add_module(
"linear",
SparseWeights(
nn.Linear(25 * cnn_out_channels[1], linear_units),
sparsity=linear_sparsity,
),
)
self.add_module("linear_bn", nn.BatchNorm1d(linear_units, affine=False))
self.add_module(
"linear_kwinner",
SampledKWinners(
percent_on=linear_percent_on,
k_inference_factor=k_inference_factor,
temperature=temperature,
eval_temperature=eval_temperature,
temperature_decay_rate=temperature_decay_rate,
relu=False,
),
)
# Classifier
self.add_module("output", nn.Linear(linear_units, 12))
self.add_module("softmax", nn.LogSoftmax(dim=1))
def __init__(self,
num_classes=10,
sparsify=False,
percent_on_fc=0.3,
percent_on_conv=0.3,
k_inference_factor=1.5,
boost_strength=1.0,
boost_strength_factor=0.9,
duty_cycle_period=1000,
hidden_units=512,
dropout=0.5,
weight_sparsity_fc=0.5,
weight_sparsity_conv=0.99,
image_size=32,
channels=3,
stats=False):
"""
Constructor for the object SimpleCNN
Args:
num_classes (int): total number of classes of the benchmark,
i.e. maximum output neurons of the model.
sparsify (bool): if we want to introduce the Kwinners and
SparseWeights layers in the model.
percent_on_fc (float): Percentage of active units in fc layers.
percent_on_conv (float): Percentage of active units in convs
layers.
k_inference_factor (float): boosting parameter. Check the
official Kwinners docs for further
details.
boost_strength (float): boosting parameter.
boost_strength_factor (float): boosting parameter.
hidden_units (int): number of units for the hidden layer.
dropout (float): dropout probability for each dropout layer.
weight_sparsity_fc (float): percentage of active weights for
each fc layer.
weight_sparsity_conv (float): percentage of active weights for
the convs layers.
image_size (int): input image size.
channels (int): number of channels in the input.
stats (bool): if we want to record sparsity statistics.
"""
super(SimpleCNN, self).__init__()
self.active_perc_list = []
self.on_idxs = [0] * hidden_units
self.hidden_units = hidden_units
self.num_classes = num_classes
self.stats = stats
if not sparsify:
self.features = nn.Sequential(
nn.Conv2d(channels, 64, kernel_size=5, stride=1, padding=0),
nn.ReLU(inplace=True),
nn.Dropout(p=dropout),
nn.Conv2d(64, 128, kernel_size=5, stride=1, padding=0),
nn.ReLU(inplace=True),
nn.Dropout(p=dropout),
nn.Flatten(),
nn.Linear(128 * 24 * 24, hidden_units),
nn.ReLU(inplace=True),
nn.Dropout(dropout),
)
self.classifier = nn.Sequential(
nn.Linear(hidden_units, num_classes))
else:
self.features = nn.Sequential(
SparseWeights2d(nn.Conv2d(channels,
64,
kernel_size=5,
stride=1,
padding=0),
weight_sparsity=weight_sparsity_conv),
nn.ReLU(inplace=True),
KWinners2d(64,
percent_on_conv,
k_inference_factor,
boost_strength,
boost_strength_factor,
duty_cycle_period,
local=True),
nn.Dropout(dropout),
SparseWeights2d(nn.Conv2d(64,
128,
kernel_size=5,
stride=1,
padding=0),
weight_sparsity=weight_sparsity_conv),
nn.ReLU(inplace=True),
KWinners2d(128,
percent_on_conv,
k_inference_factor,
boost_strength,
boost_strength_factor,
duty_cycle_period,
local=True),
nn.Dropout(dropout),
nn.Flatten(),
SparseWeights(nn.Linear(128 * 24 * 24, hidden_units),
weight_sparsity=weight_sparsity_fc),
nn.ReLU(inplace=True),
KWinners(hidden_units, percent_on_fc, k_inference_factor,
boost_strength, boost_strength_factor,
duty_cycle_period),
nn.Dropout(dropout),
)
self.classifier = nn.Sequential(
nn.Linear(hidden_units, num_classes))
def __init__(
self,
cnn_out_channels=(32, 64),
cnn_percent_on=(0.1, 0.2),
cnn_weight_sparsity=None,
linear_units=700,
linear_percent_on=0.2,
linear_weight_sparsity=None,
boost_strength=1.5,
boost_strength_factor=0.85,
k_inference_factor=1.0,
duty_cycle_period=1000,
kwinner_local=False,
cnn_sparsity=(0.4, 0.55),
linear_sparsity=0.8,
):
if cnn_weight_sparsity is not None:
warnings.warn(
"Parameter `cnn_weight_sparsity` is deprecated. Use "
"`cnn_sparsity` instead.",
DeprecationWarning,
)
cnn_sparsity = (1.0 - cnn_weight_sparsity[0],
1.0 - cnn_weight_sparsity[1])
if linear_weight_sparsity is not None:
warnings.warn(
"Parameter `linear_weight_sparsity` is deprecated. Use "
"`linear_sparsity` instead.",
DeprecationWarning,
)
linear_sparsity = 1.0 - linear_weight_sparsity
super(MNISTSparseCNN, self).__init__(
OrderedDict([
# First Sparse CNN layer
(
"cnn1",
SparseWeights2d(
nn.Conv2d(1, cnn_out_channels[0], 5),
sparsity=cnn_sparsity[0],
),
),
("cnn1_maxpool", nn.MaxPool2d(2)),
(
"cnn1_kwinner",
KWinners2d(
channels=cnn_out_channels[0],
percent_on=cnn_percent_on[0],
k_inference_factor=k_inference_factor,
boost_strength=boost_strength,
boost_strength_factor=boost_strength_factor,
duty_cycle_period=duty_cycle_period,
local=kwinner_local,
),
),
# Second Sparse CNN layer
(
"cnn2",
SparseWeights2d(
nn.Conv2d(cnn_out_channels[0], cnn_out_channels[1], 5),
sparsity=cnn_sparsity[1],
),
),
("cnn2_maxpool", nn.MaxPool2d(2)),
(
"cnn2_kwinner",
KWinners2d(
channels=cnn_out_channels[1],
percent_on=cnn_percent_on[1],
k_inference_factor=k_inference_factor,
boost_strength=boost_strength,
boost_strength_factor=boost_strength_factor,
duty_cycle_period=duty_cycle_period,
local=kwinner_local,
),
),
("flatten", Flatten()),
# Sparse Linear layer
(
"linear",
SparseWeights(
nn.Linear(16 * cnn_out_channels[1], linear_units),
sparsity=linear_sparsity,
),
),
(
"linear_kwinner",
KWinners(
n=linear_units,
percent_on=linear_percent_on,
k_inference_factor=k_inference_factor,
boost_strength=boost_strength,
boost_strength_factor=boost_strength_factor,
duty_cycle_period=duty_cycle_period,
),
),
# Classifier
("output", nn.Linear(linear_units, 10)),
("softmax", nn.LogSoftmax(dim=1)),
]))
