Model Pruning (Lottery Ticket Hypothesis) Not Reinitializing Weights

[mune98]

I am experimenting with model pruning in pytorch lightning. I noticed when pruning using the lottery ticket hypothesis (LTH), weights are not reset to the original initialization as proposed in the LTH paper and mentioned in the lightning docs. I reproduced the behaviour I faced in my own work using Tutorial 5: Transformers and Multi-Head Attention. Sharing this here to verify my analysis.

Code I added:

class CheckPruningWeight(pl.Callback):
  
    def on_train_epoch_start(self, trainer, pl_module):
        print(f"\nEPOCH {trainer.current_epoch} STARTING")
        for name in trainer.model.state_dict():
          if 'input_net.1' in name:
              print(name, trainer.model.state_dict()[name][:1])
        print('\n')

    def on_train_epoch_end(self, trainer, pl_module):
        print(f"\nEPOCH {trainer.current_epoch} ENDING")
        for name in trainer.model.state_dict():
          if 'input_net.1' in name:
              print(name, trainer.model.state_dict()[name][:1])

Pruning callback:

 pruning_callback = pl.callbacks.ModelPruning(
            pruning_fn="l1_unstructured",
            amount=0.2,
            use_global_unstructured=True,
            use_lottery_ticket_hypothesis=True,
            verbose=1,
            parameter_names=['weight'],
            resample_parameters=False,
            prune_on_train_epoch_end=False,
        )

Trainer:

 trainer = pl.Trainer(
        default_root_dir=root_dir,
        callbacks=[ModelCheckpoint(save_weights_only=True, mode="max", monitor="val_acc"),
                   pruning_callback,
                   CheckPruningWeight()],
        gpus=1 if str(device).startswith("cuda") else 0,
        max_epochs=5,
        gradient_clip_val=5,
        progress_bar_refresh_rate=1,
    )

Training Snippet:

 
  | Name                | Type               | Params
-----------------------------------------------------------
0 | input_net           | Sequential         | 352   
1 | positional_encoding | PositionalEncoding | 0     
2 | transformer         | TransformerEncoder | 8.5 K 
3 | output_net          | Sequential         | 1.4 K 
-----------------------------------------------------------
10.3 K    Trainable params
0         Non-trainable params
10.3 K    Total params
0.041     Total estimated model params size (MB)
Epoch 4: 10%
40/398 [00:00<00:08, 41.01it/s, loss=0.294, v_num=15]

EPOCH 0 STARTING
input_net.1.weight tensor([[ 0.0505, -0.1712,  0.3061, -0.2354,  0.3003, -0.0694,  0.2055, -0.0367,
         -0.1696,  0.2973]])
input_net.1.bias tensor([0.1139])


Applied `L1Unstructured`. Pruned: 0/10346 (0.00%) -> 1990/10346 (19.23%)

EPOCH 0 ENDING
input_net.1.bias tensor([0.1109])
input_net.1.weight_orig tensor([[ 0.0014, -0.0814,  0.2823, -0.3069,  0.2731, -0.1023,  0.2084, -0.0413,
         -0.2384,  0.4107]])
input_net.1.weight_mask tensor([[0., 1., 1., 1., 1., 1., 1., 1., 1., 1.]])

EPOCH 1 STARTING
input_net.1.bias tensor([0.1109])
input_net.1.weight_orig tensor([[ 0.0014, -0.0814,  0.2823, -0.3069,  0.2731, -0.1023,  0.2084, -0.0413,
         -0.2384,  0.4107]])
input_net.1.weight_mask tensor([[0., 1., 1., 1., 1., 1., 1., 1., 1., 1.]])


Applied `L1Unstructured`. Pruned: 1990/10346 (19.23%) -> 3582/10346 (34.62%)

EPOCH 1 ENDING
input_net.1.bias tensor([0.1116])
input_net.1.weight_orig tensor([[ 0.0038, -0.0283,  0.1326, -0.1825,  0.2089, -0.0469,  0.0945, -0.0116,
         -0.1252,  0.2888]])
input_net.1.weight_mask tensor([[0., 0., 1., 1., 1., 0., 1., 0., 1., 1.]])

EPOCH 2 STARTING
input_net.1.bias tensor([0.1116])
input_net.1.weight_orig tensor([[ 0.0038, -0.0283,  0.1326, -0.1825,  0.2089, -0.0469,  0.0945, -0.0116,
         -0.1252,  0.2888]])
input_net.1.weight_mask tensor([[0., 0., 1., 1., 1., 0., 1., 0., 1., 1.]])


Applied `L1Unstructured`. Pruned: 3582/10346 (34.62%) -> 4856/10346 (46.94%)

EPOCH 2 ENDING
input_net.1.bias tensor([0.1114])
input_net.1.weight_orig tensor([[ 0.0038, -0.0274,  0.1090, -0.1602,  0.2130, -0.0471,  0.0802, -0.0114,
         -0.1096,  0.2959]])
input_net.1.weight_mask tensor([[0., 0., 1., 1., 1., 0., 0., 0., 1., 1.]])

EPOCH 3 STARTING
input_net.1.bias tensor([0.1114])
input_net.1.weight_orig tensor([[ 0.0038, -0.0274,  0.1090, -0.1602,  0.2130, -0.0471,  0.0802, -0.0114,
         -0.1096,  0.2959]])
input_net.1.weight_mask tensor([[0., 0., 1., 1., 1., 0., 0., 0., 1., 1.]])

The training snippet shows the original weight tensor changing at the start of each epoch, changing values to that of the end of the previous epoch. By right, LTH should revert weights start of every epoch to the original initialization (start of epoch 0). Thoughts or correction is much appreciated.

[mune98]

As I began implementing and comparing with my own custom pruning function, I came to the realization that weight_orig is the updated weight at the end of each epoch, not the original initialized weight. Printing the actual weight at the start of each epoch gives the original initialized weight. Had a different idea of what weight_orig meant, my bad!