cuda

app.py

@@ -12,6 +12,8 @@ import annoy
 
 CONFIG_PATH = "app.ini"
 
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
 
 def parse_args():
     parser = argparse.ArgumentParser()
@@ -73,15 +75,15 @@ filenames = data["filenames"]
 
 urls = [base_url + filename for filename in filenames]
 
-model, preprocess = clip.load('RN50')
+model, preprocess = clip.load('RN50', device=device)
 
 
 def embed_text(text):
-    tokens = clip.tokenize([text])
+    tokens = clip.tokenize([text]).to(device)
     with torch.no_grad():
         text_features = model.encode_text(tokens)
     assert text_features.shape == (1, d)
-    text_features = text_features.numpy()[0]
+    text_features = text_features.cpu().numpy()[0]
     text_features /= np.linalg.norm(text_features)
     return text_features
 
@@ -95,7 +97,6 @@ def image_retrieval_from_text(text):
 
 def image_retrieval_from_image(state, selected_locally):
     selected = state[int(selected_locally)]
-    image_vector = image_features[selected][None, :]
     indices = annoy_index.get_nns_by_item(selected, n=20)
     top_urls = np.array(urls)[indices]
     return top_urls.tolist(), indices

create_embeddings.py

@@ -8,14 +8,18 @@ import pickle
 
 
 def do_batch(batch, embeddings):
-    image_batch = torch.tensor(np.stack(batch))
+    image_batch = torch.tensor(np.stack(batch)).to(device)
     with torch.no_grad():
         image_features = model.encode_image(image_batch).float()
-        embeddings += image_features.numpy().tolist()
+        embeddings += image_features.cpu().numpy().tolist()
         print(f"{len(embeddings)} done")
 
 
-model, preprocess = clip.load('RN50')
+# even though it's not worth bothering with cuda,
+# because 98% of the run time is preprocessing on the cpu.
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+model, preprocess = clip.load('RN50', device=device)
 
 limit = 1e9
 batch_size = 100