INNER CODE UNIT · Python

set_seed

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

def set_seed(seed):
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    random.seed(seed)
    torch.backends.cudnn.deterministic = True


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

View source record →