Method predict

core/defenses/AutoEncoderDefense.py:216–252

Apply AutoEncoder defense method to input data and get the predicts. Args: model (torch.nn.Module): Network. data (torch.Tensor): Input data (between 0.0 and 1.0), shape: (N, C, H, W), dtype: torch.float32. schedule (dict): Schedule for predicting.

(self, model, data, schedule)

            return predict_digits

    def predict(self, model, data, schedule):
        """Apply AutoEncoder defense method to input data and get the predicts.

        Args:
            model (torch.nn.Module): Network.
            data (torch.Tensor): Input data (between 0.0 and 1.0), shape: (N, C, H, W), dtype: torch.float32.
            schedule (dict): Schedule for predicting.

        Returns:
            torch.Tensor: The predicts.
        """
        preprocessed_data = self.preprocess(data)

        if 'test_model' in schedule:
            model.load_state_dict(torch.load(schedule['test_model']), strict=False)

        # Use GPU
        if 'device' in schedule and schedule['device'] == 'GPU':
            if 'CUDA_VISIBLE_DEVICES' in schedule:
                os.environ['CUDA_VISIBLE_DEVICES'] = schedule['CUDA_VISIBLE_DEVICES']

            assert torch.cuda.device_count() > 0, 'This machine has no cuda devices!'
            assert schedule['GPU_num'] >0, 'GPU_num should be a positive integer'
            print(f"This machine has {torch.cuda.device_count()} cuda devices, and use {schedule['GPU_num']} of them to train.")

            if schedule['GPU_num'] == 1:
                device = torch.device("cuda:0")
            else:
                gpus = list(range(schedule['GPU_num']))
                model = nn.DataParallel(model.cuda(), device_ids=gpus, output_device=gpus[0])
                # TODO: DDP training
                pass
        # Use CPU
        else:
            device = torch.device("cpu")

        return self._predict(model, preprocessed_data, device, schedule['batch_size'], schedule['num_workers'])

    def test(self, model, dataset, schedule):
        """Test AutoEncoder on dataset.