Resulting sparsity is 0 for all the layers about sparseml HOT 5 CLOSED

train_batch_size = 64
val_batch_size = 1

shape = (224, 224)

train_loader = DataLoader(
    train_dset,
    batch_size=train_batch_size,
    shuffle=True,
    num_workers=6)

lr = 0.0005

model = smp.Unet('efficientnet-b0', activation=None, encoder_weights=None)
model = model.cuda()

optim = torch.optim.Adam(model.parameters(), lr=lr)

epochs = 50
no_batches_val = math.ceil(val_size / val_batch_size)
loss = BCEDiceLoss()
no_batches_train = math.ceil(train_size / train_batch_size)

checkpoint = torch.load('efbn0_224_sparsify_test(14).tar')
model.load_state_dict(checkpoint['model_state_dict'])

manager = ScheduledModifierManager.from_yaml('sparsify_recipe.yaml')
optimizer = manager.modify(model, optim, steps_per_epoch=len(train_loader))
train_loss_history = []
train_iou_history = []
for epoch in tqdm(range(manager.max_epochs)):

    model = model.cuda().train()

    for (imgs, labels, _) in tqdm(train_loader):
        optim.zero_grad()

        imgs, labels = imgs.cuda(), labels.cuda()

        out = model(imgs)
        batch_train_loss = loss(out, labels)

        batch_train_loss.backward()
        optim.step()
    
manager.finalize(model)

torch.save({
            'epoch': epoch,
            'model_state_dict': model.state_dict(),
            },  'efbn0_224_sparse_final(' + str(epoch) + ').tar')

import torch
from torch import nn

def f_score(pr, gt, beta=1, eps=1e-7, threshold=None, activation='sigmoid'):
    """
    Args:
        pr (torch.Tensor): A list of predicted elements
        gt (torch.Tensor):  A list of elements that are to be predicted
        beta (float): positive constant
        eps (float): epsilon to avoid zero division
        threshold: threshold for outputs binarization
    Returns:
        float: F score
    """

    if activation is None or activation == "none":
        activation_fn = lambda x: x
    elif activation == "sigmoid":
        activation_fn = torch.nn.Sigmoid()
    elif activation == "softmax2d":
        activation_fn = torch.nn.Softmax2d()
    else:
        raise NotImplementedError(
            "Activation implemented for sigmoid and softmax2d"
        )

    pr = activation_fn(pr)

    if threshold is not None:
        pr = (pr > threshold).float()


    tp = torch.sum(gt * pr)
    fp = torch.sum(pr) - tp
    fn = torch.sum(gt) - tp

    score = ((1 + beta ** 2) * tp + eps) \
            / ((1 + beta ** 2) * tp + beta ** 2 * fn + fp + eps)

class DiceLoss(nn.Module):
    __name__ = 'dice_loss'

    def __init__(self, eps=1e-7, activation='sigmoid'):
        super().__init__()
        self.activation = activation
        self.eps = eps

    def forward(self, y_pr, y_gt):
        return 1 - f_score(y_pr, y_gt, beta=1., eps=self.eps, threshold=None, activation=self.activation)

class BCEDiceLoss(DiceLoss):
    __name__ = 'bce_dice_loss'

    def __init__(self, eps=1e-7, activation='none'):
        super().__init__(eps, activation)
        self.bce = nn.BCEWithLogitsLoss(reduction='mean')

    def forward(self, y_pr, y_gt):
        dice = super().forward(y_pr, y_gt)
        bce = self.bce(y_pr, y_gt)
        return dice + bce