INNER CODE UNIT · Python

count_parameters

haitongli/knowledge-distillation-pytorch · count_model_size.py:21

def count_parameters(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)


if __name__ == '__main__':

    model_size = 0

    args = parser.parse_args()
    cnn_dir = 'experiments/cnn_distill'
    json_path = os.path.join(cnn_dir, 'params.json')
    assert os.path.isfile(json_path), "No json configuration file found at {}".format(json_path)
    params = utils.Params(json_path)

    if args.model == "resnet18":
        model = resnet.ResNet18()
        model_checkpoint = 'experiments/base_resnet18/best.pth.tar'

View source record →