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)), ]))