ljzycmd/VDTR

About Computational Complexity

wgp666 opened this issue · 1 comments

Hello! Can you share how you calculate the GFLOPs and Params. I used thop to calculate and got your model GFLOPs = 1362.78 and Params = 16.01M.

The thop package may be inaccurate due to the customized network structure. I use the flop_count function here to calculate the model complexity. As for the Params, I sum up all the parameters in the network.

def calculate_model_param(model: nn.Module):
    total_params = 0
    trainable_params = 0
    non_trainable_params = 0

    # traverse model.parameters()
    for param in model.parameters():
        mul_value = np.prod(param.size()) 
        total_params += mul_value
        if param.requires_grad:
            trainable_params += mul_value  # trainable
        else:
            non_trainable_params += mul_value  # non-trainable

    total_params /= 1e6
    trainable_params /= 1e6
    non_trainable_params /= 1e6

    print(f'Total params: {total_params} M.')
    print(f'Trainable params: {trainable_params} M.')
    print(f'Non-trainable params: {non_trainable_params} M.')

Hope this can help you!