Molbap HF Staff commited on
Commit
b702ae2
·
1 Parent(s): 2ccc01b

gpu optims

Browse files
Files changed (1) hide show
  1. modular_graph_and_candidates.py +4 -3
modular_graph_and_candidates.py CHANGED
@@ -130,7 +130,7 @@ def embedding_similarity_clusters(models_root: Path, missing: List[str], thr: fl
130
  all_embeddings = []
131
 
132
  print(f"Encoding embeddings for {len(names)} models...")
133
- batch_size = 8
134
 
135
  for i in tqdm(range(0, len(names), batch_size), desc="Batches", leave=False):
136
  batch_names = names[i:i+batch_size]
@@ -142,9 +142,10 @@ def embedding_similarity_clusters(models_root: Path, missing: List[str], thr: fl
142
  all_embeddings.append(emb)
143
  print(f"✓ Completed batch of {len(batch_names)} models")
144
 
145
- # Clear GPU cache every 5 batches to prevent memory accumulation
146
- if i % (5 * batch_size) == 0 and torch.cuda.is_available():
147
  torch.cuda.empty_cache()
 
148
  print(f"🧹 Cleared GPU cache after batch {i//batch_size + 1}")
149
 
150
  except Exception as e:
 
130
  all_embeddings = []
131
 
132
  print(f"Encoding embeddings for {len(names)} models...")
133
+ batch_size = 4 # Reduced to be more conservative
134
 
135
  for i in tqdm(range(0, len(names), batch_size), desc="Batches", leave=False):
136
  batch_names = names[i:i+batch_size]
 
142
  all_embeddings.append(emb)
143
  print(f"✓ Completed batch of {len(batch_names)} models")
144
 
145
+ # Clear GPU cache every 3 batches to prevent memory accumulation
146
+ if i % (3 * batch_size) == 0 and torch.cuda.is_available():
147
  torch.cuda.empty_cache()
148
+ torch.cuda.synchronize() # Force GPU sync
149
  print(f"🧹 Cleared GPU cache after batch {i//batch_size + 1}")
150
 
151
  except Exception as e: