Modifying gradients stop training

grid_world: <p>I am trying to train a pruned neural network which has (say) 36% sparsity which means that 36% of the trainable parameters are 0s. I am using a simple LeNet-300-100 dense network for this:</p>
<pre><code class="lang-auto">class LeNet300(nn.Module):
    def __init__(self):
        super().__init__()
        
        # Define layers-
        self.fc1 = nn.Linear(in_features = 28 * 28 * 1, out_features = 300)
        self.fc2 = nn.Linear(in_features = 300, out_features = 100)
        self.output = nn.Linear(in_features = 100, out_features = 10)
        
        # self.weights_initialization()
    
    
    def forward(self, x):
        out = F.leaky_relu(self.fc1(x))
        out = F.leaky_relu(self.fc2(out))
        return self.output(out)
</code></pre>
<p>After 36% sparsity, the number of trainable parameters = 213288. Originally, without pruning, it has 266610 non-zero parameters.</p>
<p>The layer-wise sparsity looks as:</p>
<pre><code class="lang-auto">for param in model.parameters():
    print(f"{param.size()} has {torch.count_nonzero(param)} surviving weights")

'''
torch.Size([300, 784]) has 146615 surviving weights
torch.Size([300]) has 300 surviving weights
torch.Size([100, 300]) has 22743 surviving weights
torch.Size([100]) has 100 surviving weights
torch.Size([10, 100]) has 862 surviving weights
torch.Size([10]) has 10 surviving weights
'''
</code></pre>
<p>The code I am using to preserve sparsity is:</p>
<p>In .grad (<code>gradient_t</code>), replace computed gradients with 0s in place positions (within the matrix) where the trainable parameters (<code>wts</code>) are pruned:</p>
<pre><code class="lang-auto"># Compute loss-
J = loss(outputs, labels)
            
# Empty accumulated gradients-
optimizer.zero_grad()
            
# Perform backprop-
J.backward()
            
for name, param in model.named_parameters():
    wts = param.data.clone().detach()
    gradient_t = param.grad
    gradient_t = torch.where(wts == 0., 0., gradient_t)
    param.grad = gradient_t
        
# Update parameters-
optimizer.step()
</code></pre>
<p>However, after this, the training basically freezes, and the loss and accuracy stop changing.</p>
<p>?</p>

KFrank: <p>Hi Arjun!</p>
<aside class="quote no-group" data-username="grid_world" data-post="1" data-topic="186036" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/grid_world/48/33204_2.png" class="avatar"> grid_world:</div>
<blockquote>
<pre><code class="lang-auto">for name, param in model.named_parameters():
    wts = param.data.clone().detach()
    gradient_t = param.grad
    gradient_t = torch.where(wts == 0., 0., gradient_t)
    param.grad = gradient_t
        
# Update parameters-
optimizer.step()
</code></pre>
</blockquote>
</aside>
<p>In principle your scheme should work.  Consider this example script:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
print (torch.__version__)

s = torch.tensor ([2.0, 0.0, -1.0], requires_grad = True)   # a trainable parameter
t = torch.tensor ([1.0, 1.0, 1.0])

optimizer = torch.optim.SGD ([s], lr = 0.2)                 # plain-vanilla SGD

for  i in range (10):
    loss = ((s - t)**2).sum()
    optimizer.zero_grad()
    loss.backward()
    wts = s.data.clone().detach()                           # don't use .data -- it is deprecated
    gradient_t = s.grad
    gradient_t = torch.where (wts == 0., 0., gradient_t)
    s.grad = gradient_t
    optimizer.step()
    print (s)                                               # s trains nicely with s[1] frozen
</code></pre>
<p>Here is its output:</p>
<pre><code class="lang-plaintext">.0.1
tensor([ 1.6000,  0.0000, -0.2000], requires_grad=True)
tensor([1.3600, 0.0000, 0.2800], requires_grad=True)
tensor([1.2160, 0.0000, 0.5680], requires_grad=True)
tensor([1.1296, 0.0000, 0.7408], requires_grad=True)
tensor([1.0778, 0.0000, 0.8445], requires_grad=True)
tensor([1.0467, 0.0000, 0.9067], requires_grad=True)
tensor([1.0280, 0.0000, 0.9440], requires_grad=True)
tensor([1.0168, 0.0000, 0.9664], requires_grad=True)
tensor([1.0101, 0.0000, 0.9798], requires_grad=True)
tensor([1.0060, 0.0000, 0.9879], requires_grad=True)
</code></pre>
<aside class="quote no-group">
<blockquote>
<p>However, after this,</p>
</blockquote>
</aside>
<p>After what, specifically?</p>
<aside class="quote no-group">
<blockquote>
<p>the training basically freezes,</p>
</blockquote>
</aside>
<p>What does “basically freezes” mean.  Do your weights stop changing<br>
<em>entirely</em> or does your training just slow down some?</p>
<aside class="quote no-group">
<blockquote>
<p>and the loss and accuracy stop changing.</p>
</blockquote>
</aside>
<p>Do you get any non-zero gradients?  When you call <code>optimizer.step()</code><br>
do the weights with non-zero gradients change at all?  Do your <code>outputs</code><br>
(predictions) change at all?  If your predictions change some, do your<br>
loss and accuracy change at least a little bit, or do then not change at<br>
all?</p>
<p>As an aside, although it isn’t the cause of your problem, you shouldn’t<br>
use <code>.data</code>, as it’s deprecated an can break things.  It also isn’t necessary<br>
in your case, but I would suggest modifying the gradients that you want to<br>
zero out inplace, and wrapping that whole bit in a <code>with torch.no_grad():</code><br>
block.</p>
<p>Best.</p>
<p>K. Frank</p>

grid_world: <p>Hey Frank,</p>
<p>Thanks for your detailed answer. Can you provide an example of modifying the gradients by 0ing them out inplace?</p>

grid_world: <p><a class="mention" href="/u/kfrank">@KFrank</a> I have uploaded an example reproducible code as a jupyter notebook which you can access <a href="https://github.com/arjun-majumdar/Neural_Network_Pruning/blob/main/Pruning_Example_torch.ipynb" rel="noopener nofollow ugc">here</a>.</p>
<p>The code to implement global, unstructured, magnitude pruning is:</p>
<pre><code class="lang-auto">def prune_globally(model, pruning_percentile = 20):
    # Python 3 list to hold layer-wise weights-
    pruned_weights = []
    
    for param in model.parameters():
        wts = np.copy(param.detach().cpu().numpy())
        pruned_weights.append(wts)
    
    del param, wts
    
    # Flatten all numpy arrays-
    pruned_weights_flattened = [layer.flatten() for layer in pruned_weights]

    threshold = np.percentile(a = abs(np.concatenate(pruned_weights_flattened)), q = pruning_percentile)
    # print("\nFor p = {0:.2f}% of weights to be pruned, threshold = {1:.4f}\n".format(p, threshold))
    
    # Prune conv and dense layers-
    # bias and batch-norm is NOT pruned.
    for layer in pruned_weights:
        if len(layer.shape) == 4:
            layer[abs(layer) &lt; threshold] = 0
        elif len(layer.shape) == 2:
            layer[abs(layer) &lt; threshold] = 0
    
    
    i = 0
    model_d = dict()

    for name, params in model.named_parameters():
        if pruned_weights[i].shape == params.shape:
            model_d[name] = torch.from_numpy(pruned_weights[i])

        i += 1
        
        
    state_d = model.state_dict()

    for layer_name in model_d:
        # if pruned_model.state_dict().get(layer_name) is not None:
        if state_d.get(layer_name) is not None:
            # print(layer_name)
            state_d[layer_name] = model_d.get(layer_name)

    model.load_state_dict(state_d)
        
    return None
</code></pre>
<p>Within the notebook, after pruning the trained model for 20% sparsity in cell number 35 with:</p>
<p><code>prune_globally(model, pruning_percentile = sparsity_percentage[0])</code></p>
<p>The number of trainable surviving parameters = 213288. Upon re-training the pruned model in cell number 37 with:</p>
<pre><code class="lang-auto"># Train until convergence pruned initialized model-
history_pr = train_until_convergence(
    model = model, train_loader = train_loader,
    test_loader = test_loader, train_dataset = train_dataset,
    test_dataset = test_dataset, num_epochs = 20
)
</code></pre>
<p>The sparsity is destroyed as the number of trainable parameters falls back to 266610 and not 213288.</p>
<p>Something is bringing back the parameters which are supposed to stay pruned at 0??</p>

KFrank: <p>Hi Arjun!</p>
<aside class="quote no-group" data-username="grid_world" data-post="3" data-topic="186036" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/grid_world/48/33204_2.png" class="avatar"> grid_world:</div>
<blockquote>
<p>Can you provide an example of modifying the gradients by 0ing them out inplace?</p>
</blockquote>
</aside>
<p>In general, I would look to <code>.copy_()</code>, but in this case I might use <code>.mul_()</code>:</p>
<pre data-code-wrap="python"><code class="lang-python">with  torch.no_grad():
    for  name, param in model.named_parameters():
        param.grad.mul_ (~(param == 0.0))
</code></pre>
<p>Best.</p>
<p>K. Frank</p>

grid_world: <p>I changed the train function to include your code:</p>
<pre><code class="lang-auto">def train_model_progress(
    model, 
    train_loader, train_dataset
):
    '''
    Function to perform one epoch of training by using 'train_loader'.
    Returns loss and number of correct predictions for this epoch.
    '''
    running_loss = 0.0
    running_corrects = 0.0
    
    model.train()
    
    with tqdm(train_loader, unit = 'batch') as tepoch:
        for images, labels in tepoch:
            tepoch.set_description(f"Training: ")
            
            images = images.view(-1, 28 * 28 * 1)
            
            images = images.to(device)
            labels = labels.to(device)
            
            # Get model predictions-
            outputs = model(images)
            
            # Compute loss-
            J = loss(outputs, labels)
            
            # Empty accumulated gradients-
            optimizer.zero_grad()
            
            # Perform backprop-
            J.backward()
            
            
            with  torch.no_grad():
                for  name, param in model.named_parameters():
                    param.grad.mul_(~(param == 0.0))
            
            # Update parameters-
            optimizer.step()
                        
            # Compute model's performance statistics-
            running_loss += J.item() * images.size(0)
            _, predicted = torch.max(outputs, 1)
            running_corrects += torch.sum(predicted == labels.data)
            
            tepoch.set_postfix(
                loss = running_loss / len(train_dataset),
                accuracy = (running_corrects.double().cpu().numpy() / len(train_dataset)) * 100
            )
            
    
    train_loss = running_loss / len(train_dataset)
    train_acc = (running_corrects.double() / len(train_dataset)) * 100

    return train_loss, train_acc.cpu().numpy()
</code></pre>
<p>But after the first pruning round to remove 20% of parameters, the number of trainable parameters falls back to the original number of parameters, therefore destroying pruning - Alas!</p>

grid_world: <p><a class="mention" href="/u/kfrank">@KFrank</a> I think <code>optimizer.step()</code> destroys sparsity. Look at this minimal code:</p>
<pre><code class="lang-auto">images, labels = next(iter(train_loader))
images = images.view(-1, 28 * 28 * 1)
images = images.to(device)
labels = labels.to(device)
            
# Get model predictions-
outputs = model(images)
            
# Compute loss-
J = loss(outputs, labels)
            
# Empty accumulated gradients-
optimizer.zero_grad()
            
# Perform backprop-
J.backward()

# before grad multiplication-
for param in model.parameters():
    print(f"{param.size()}; wts = {torch.count_nonzero(param)}",
          f" &amp; grads = {torch.count_nonzero(param.grad)}"
         )
'''
torch.Size([300, 784]); wts = 185620  &amp; grads = 235200
torch.Size([300]); wts = 300  &amp; grads = 300
torch.Size([100, 300]); wts = 26318  &amp; grads = 30000
torch.Size([100]); wts = 100  &amp; grads = 100
torch.Size([10, 100]); wts = 940  &amp; grads = 1000
torch.Size([10]); wts = 10  &amp; grads = 10
'''

# after grad multiplication-
with  torch.no_grad():
    # for name, param in model.named_parameters():
    for param in model.parameters():
        param.grad.mul_(~(param == 0.0))
        
        print(f"{param.size()}; wts = {torch.count_nonzero(param)}",
          f" &amp; grads = {torch.count_nonzero(param.grad)}"
         )
'''
torch.Size([300, 784]); wts = 185620  &amp; grads = 185620
torch.Size([300]); wts = 300  &amp; grads = 300
torch.Size([100, 300]); wts = 26318  &amp; grads = 26318
torch.Size([100]); wts = 100  &amp; grads = 100
torch.Size([10, 100]); wts = 940  &amp; grads = 940
torch.Size([10]); wts = 10  &amp; grads = 10
'''

count_surviving_params(model)
# 213288

optimizer.step()

count_surviving_params(model)
# 266610
</code></pre>
<p>Just after <code>optimizer.step()</code> is called, the pruned parameters seem to be back and tosses sparsity out of the window?!?</p>

KFrank: <p>Hi Arjun!</p>
<aside class="quote no-group quote-modified" data-username="grid_world" data-post="7" data-topic="186036" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/grid_world/48/33204_2.png" class="avatar"> grid_world:</div>
<blockquote>
<p>I think <code>optimizer.step()</code> destroys sparsity.<br>
…<br>
Just after <code>optimizer.step()</code> is called, the pruned parameters seem to be back and tosses sparsity out of the window?!?</p>
</blockquote>
</aside>
<p>You haven’t told us which optimizer you are using, but some optimizers,<br>
even <code>SGD</code> with non-zero momentum, will modify parameters whose gradients<br>
are (or have been forced to be) zero.</p>
<p>My post above shows a simple example where your scheme works with<br>
plain-vanilla <code>SGD</code>.</p>
<p>You have two choices:  Zero out the gradients and use an optimizer that<br>
does not modify parameters that have zero gradients; or zero out the<br>
parameters themselves after <code>optimizer.step()</code>.</p>
<p>Best.</p>
<p>K. Frank</p>

grid_world: <p>In Lottery Ticket Hypothesis and other similar papers, the authors use Adam and/or SGD optimizer.</p>
<p>The code that I had implemented in TensorFlow 2.x is:</p>
<pre><code class="lang-auto"># Define 'train_one_step()' and 'test_step()' functions here-
    @tf.function
    def train_one_step(model, mask_model, optimizer, x, y):
        '''
        Function to compute one step of gradient descent optimization
        '''
        with tf.GradientTape() as tape:
            # Make predictions using defined model-
            y_pred = model(x)

            # Compute loss-
            loss = loss_fn(y, y_pred)

        # Compute gradients wrt defined loss and weights and biases-
        grads = tape.gradient(loss, model.trainable_variables)

        # type(grads)c
        # list

        # List to hold element-wise multiplication between-
        # computed gradient and masks-
        grad_mask_mul = []

        # Perform element-wise multiplication between computed gradients and masks-
        for grad_layer, mask in zip(grads, mask_model.trainable_weights):
            grad_mask_mul.append(tf.math.multiply(grad_layer, mask))

        # Apply computed gradients to model's weights and biases-
        optimizer.apply_gradients(zip(grad_mask_mul, model.trainable_variables))

        # Compute accuracy-
        train_loss(loss)
        train_accuracy(y, y_pred)

        return None
</code></pre>
<p><code>mask_model</code> is a binary mask of 0/1 denoting which parameter has been pruned, and which is not. Here, the sparsity is maintained with <code>grad_mask_mul</code> (computed as element-wise multiplication between gradients and binary mask, and:</p>
<pre><code class="lang-auto">optimizer.apply_gradients(zip(grad_mask_mul, model.trainable_variables))
</code></pre>
<p>due to multiplication of gradients with the binary mask (stored in <code>grad_mask_mul</code>) which is then applied to model’s trainable parameters.</p>
<p>This worked for SGD and Adam.</p>
<p>Is there something similar in PyTorch ?</p>