amazon-science / patchcore-inspection

Apache License 2.0
779 stars 155 forks source link

Assign GPU id to the FaissNN search #62

Open YoojLee opened 1 year ago

YoojLee commented 1 year ago

Issue: #61

Description of changes: I added the code fraction to assign a GPU id to the GpuIndexFlatL2 by:

In src/patchcore/common.py:

class FaissNN(object):
    def __init__(self, on_gpu: bool = False, num_workers: int = 4, device: Union[int,torch.device]=0) -> None:
       ...
       if isinstance(device, torch.device):
                  device = int(torch.cuda.current_device())
              self.device = device
       ...

    def _create_index(self, dimension):
        if self.on_gpu:
            gpu_config = faiss.GpuIndexFlatConfig()
            gpu_config.device = self.device
            return faiss.GpuIndexFlatL2(
                faiss.StandardGpuResources(), dimension, gpu_config
            )
        return faiss.IndexFlatL2(dimension)

At line 297 in run_patchcore.py:

...
def get_patchcore(input_shape, sampler, device):
        nn_method = patchcore.common.FaissNN(faiss_on_gpu, faiss_num_workers, device)
...

At line 213 and 222 In load_and_evaluate.py:

def patch_core_loader(patch_core_paths, faiss_on_gpu, faiss_num_workers):
    ...
            if n_patchcores == 1:
                nn_method = patchcore.common.FaissNN(faiss_on_gpu, faiss_num_workers, device)
                patchcore_instance = patchcore.patchcore.PatchCore(device)
                patchcore_instance.load_from_path(
                    load_path=patch_core_path, device=device, nn_method=nn_method
                )
                loaded_patchcores.append(patchcore_instance)
            else:
                for i in range(n_patchcores):
                    nn_method = patchcore.common.FaissNN(
                        faiss_on_gpu, faiss_num_workers, device
                    )

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.