<jupyter_start><jupyter_text>Set the arguments before we start<jupyter_code>args_dict = {
    "prune_ratio": 0.4,
    "sparsity_reg": "none",
    "retrain_mode": "weight-rewinding",
    
    "model_path": "./saved_models/oneshot",
    "seed": 2020,
    "batch_size": 128,
    "init_lr": 0.001,
    "weight_decay": 1e-4,
    "train_epochs": 50,
    "retrain_epochs": 50,
    "rewind_epoch": 1.4, # 1.4 chosen as per the "The Lottery Ticket Hypothesis at Scale" paper
    "use_early_stop": True,
    "patience":10
}

class Args():
    def __init__(self):
        pass
args = Args()
for k,v in args_dict.items():
    setattr(args, k, v)<jupyter_output><empty_output><jupyter_text>Get the DataLoaders and the Model<jupyter_code>dataset = DATASETS['MNIST']
train_loader, val_loader, test_loader = dataset['loaders'](args, seed=2020)
criterion = criterion = nn.CrossEntropyLoss().cuda()
model = LeNet(dataset['num_classes'], criterion, device)
model = model.to(device)
optimizer = torch.optim.SGD(model.parameters(), args.init_lr, momentum=0.9, weight_decay=1e-4)
lr_scheduler = None<jupyter_output><empty_output><jupyter_text>Define the parameters to prunecheck the parameter list by running the following cell<jupyter_code>[name for name, _ in model.named_parameters()]
def parameters_to_prune(m):
    return [
        (m.conv1, 'weight'), (m.conv2, 'weight'),
        (m.fc1, 'weight'), (m.fc2, 'weight'),
    ]<jupyter_output><empty_output><jupyter_text>Lets get the Lottery Ticketreference: [Lottery Ticket Hypothesis](https://arxiv.org/abs/1803.03635)<jupyter_code>one_shot_pruning = OneShotPruning(model, parameters_to_prune, optimizer, lr_scheduler, train_loader, val_loader, compute_accuracy, device, args)
_, ckpt_orig_path, ckpt_pruned_path, loss_history_orig, loss_history_pruned = one_shot_pruning.get_ticket()<jupyter_output><empty_output>
StarcoderdataJupyterStructuredCleanDedup