improve gpu memroy efficiency

Author: yztxwd

test/test_cav_trainer.py

@@ -279,7 +279,7 @@ def test_all(self):
         cav_trainer.set_control(builder.control_concepts[0], num_samples=100)
 
         cav_trainer.train_concepts(
-            builder.concepts, 100, output_dir="data/cavs/", num_processes=1, backend='torch'
+            builder.concepts, 100, output_dir="data/cavs/", num_processes=2, backend='torch'
         )
         cav_trainer.train_concepts(
             builder.concepts, 100, output_dir="data/cavs/", num_processes=2

tpcav/cavs.py

@@ -5,6 +5,7 @@
 
 import logging
 import os
+import gc
 from pathlib import Path
 from typing import Iterable, List, Optional, Tuple, Union
 import time
@@ -134,6 +135,10 @@ def fit(self, train_val_avs: np.ndarray, train_val_ls: np.ndarray):
             if (best_loss is None) or (loss < best_loss):
                 best_loss = loss
                 best_state_dict = state_dict
+
+            del model
+            gc.collect()
+            torch.cuda.empty_cache()
 
         self.best_model = _TorchLinear(self.input_dim, self.num_class)
         self.best_model.load_state_dict(best_state_dict)
@@ -277,6 +282,12 @@ def _eval(avs, l, name: str):
     assert len(weights.shape) == 2 and weights.shape[0] == 2
     torch.save(weights, output_dir / "classifier_weights.pt")
 
+    if backend == 'torch':
+        # release gpu memroy
+        del clf.best_model
+        gc.collect()
+        torch.cuda.empty_cache()
+
     return test_fscore, weights[0]