INNER CODE UNIT · Python

get_model

digantamisra98/Mish · PyTorch Benchmarks/train_imagenet.py:167

def get_model(arch, depth):
    assert arch == "resnet"
    model = ResidualNet("ImageNet", depth, 1000)
    print(model)
    return model


def get_model_stats(model, device, verbose=True):
    with torch.cuda.device(device):
        macs, params = get_model_complexity_info(
            model, (3, 224, 224), as_strings=False, print_per_layer_stat=False
        )
        if verbose:
            print("{:<30}  {:<8}".format("Computational complexity: ", int(macs)))
            print("{:<30}  {:<8}".format("Number of parameters: ", params))
    return macs, params

View source record →