fix trojannnn neuron selection #105

trojanvision/attacks/backdoor/trojannn.py

@@ -113,11 +113,13 @@ def get_neuron_idx(self) -> torch.Tensor:
         if weight.dim() > 2:
             weight = weight.flatten(2).mean(2)
         weight = weight.mean(0)
-        return weight.argsort(descending=False)[:self.neuron_num]
+        return weight.argsort(descending=True)[:self.neuron_num]
 
     def get_neuron_value(self, trigger_input: torch.Tensor, neuron_idx: torch.Tensor) -> float:
         r"""Get average neuron activation value of :attr:`trigger_input` for :attr:`neuron_idx`.
 
+        The feature map is obtained by calling :meth:`ImageModel.get_layer()`.
+
         Args:
             trigger_input (torch.Tensor): Triggered input tensor with shape ``(N, C, H, W)``.
             neuron_idx (torch.Tensor): Neuron index list tensor with shape ``(self.neuron_num)``.
@@ -138,6 +140,8 @@ def preprocess_mark(self, neuron_idx: torch.Tensor):
         :any:`torch.optim.lr_scheduler.CosineAnnealingLR`
         with tanh objective funcion.
 
+        The feature map is obtained by calling :meth:`ImageModel.get_layer()`.
+
         Args:
             neuron_idx (torch.Tensor): Neuron index list tensor with shape ``(self.neuron_num)``.
         """