INNER CODE UNIT · Python

get_model_stats

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

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


def get_loss_optim(model, device, lr, momentum, weight_decay):
    criterion = nn.CrossEntropyLoss().to(device)

    optimizer = torch.optim.SGD(
        model.parameters(), lr, momentum=momentum, weight_decay=weight_decay
    )
    return criterion, optimizer

View source record →