INNER CODE UNIT · Python

feature_extractor

GaParmar/clean-fid · cleanfid/features.py:22

def feature_extractor(
    name: str = "torchscript_inception",
    device=torch.device("cuda"),
    resize_inside=False,
    use_dataparallel=True,
):
    device_ids = [device]
    if name == "torchscript_inception":
        path = "./" if platform.system() == "Windows" else "/tmp"
        model = InceptionV3W(
            path,
            download=True,
            resize_inside=resize_inside,
        ).to(device)
        model.eval()
        if use_dataparallel:
            model = torch.nn.DataParallel(model, device_ids=device_ids)

View source record →