make device a property.

sustainability-lab · Oct 28, 2023 · 921690f · 921690f
1 parent d14d4aa
commit 921690f
Show file tree

Hide file tree

Showing 4 changed files with 9 additions and 3 deletions.
diff --git a/astra/torch/al/strategies/base.py b/astra/torch/al/strategies/base.py
@@ -38,6 +38,7 @@ def __init__(
             acquisitions = [acquisitions]
         self.acquisitions = {acq.__class__.__name__: acq for acq in acquisitions}
 
+    @property
     def device(self):
         return self.dummy_param.device
 

diff --git a/astra/torch/al/strategies/diversity.py b/astra/torch/al/strategies/diversity.py
@@ -47,11 +47,16 @@ def query(
             # Get the features for the context
             context_features_list = []
             for x, _ in context_data_loader:
-                context_features = net(x)
+                context_features = net(x.to(self.device))
                 context_features_list.append(context_features)
             context_features = torch.cat(context_features_list, dim=0)  # (context_dim, feature_dim)
 
             best_indices = {}
+            # TODO: Fix this for loop to do the following:
+            #    - Get the max score. Get corresponding index.
+            #    - Add that index to selected_indices and also to pool indices.
+            #    - Remove that index from pool indices.
+            #    - Repeat until len(selected_indices) == n_query_samples.
             for acq_name, acquisition in self.acquisitions.items():
                 scores = acquisition.acquire_scores(pool_features, context_features)
                 selected_indices = torch.topk(scores, n_query_samples).indices

diff --git a/astra/torch/al/strategies/ensemble.py b/astra/torch/al/strategies/ensemble.py
@@ -45,7 +45,7 @@ def query(
             for x, _ in data_loader:
                 net_logits_list = []
                 for model in net:
-                    net_logits = model(x.to(self.device()))[np.newaxis, ...]
+                    net_logits = model(x.to(self.device))[np.newaxis, ...]
                     net_logits_list.append(net_logits)
                 logits = torch.cat(net_logits_list, dim=0)  # (n_nets, batch_dim, n_classes)
                 logits_list.append(logits)

diff --git a/astra/torch/al/strategies/mc.py b/astra/torch/al/strategies/mc.py
@@ -47,7 +47,7 @@ def query(
         with torch.no_grad():
             logits_list = []
             for x, _ in data_loader:
-                vx = x[np.newaxis, ...].repeat(*repeats).to(self.device())
+                vx = x[np.newaxis, ...].repeat(*repeats).to(self.device)
                 logits = torch.vmap(net, randomness="different")(vx)
                 logits_list.append(logits)
             logits = torch.cat(logits_list, dim=1)  # (n_mc_samples, pool_dim, n_classes)