mhamilton723 / STEGO

Unsupervised Semantic Segmentation by Distilling Feature Correspondences
MIT License
711 stars 142 forks source link

What are the minimal hardware requirements? #53

Closed codingbutstillalive closed 1 year ago

codingbutstillalive commented 1 year ago

I would like to set up a minimal test case and compute segmentations for a few images on my 2018 Macbook (without a Nvidia GPU). Is it worth trying, or will I need other hardware resources?

The paper says that the Segmentation head on top is very lightweight.

atthom commented 1 year ago

I would like to know that too.

Running on a RTX 3050 Ti Laptop I can pass the knn_precompute.py if I decrease the batch_size to 8. But I fail to run train_segmentation, It seems I need more than 16Go of memory for my trainingset.

mhamilton723 commented 1 year ago

I ran on a NVIDIA V-100 (if memory serves me correctly). You can adjust the batch sizes so things fit on your GPUs. The segmentation head is a small nonlinear mode, so the total cost of the evaluation is not much more than DINO itself. Training is a different story, and to get reasonable batch sizes you need a reasonable GPU or multiple GPUs.