code: def test_knn_fsdp(self) -> None:
      with in_temporary_directory() as pretrain_dir:
          # Run a pre-training to have some weights to being with
          pretrain_config = self._create_pretraining_config(
              with_fsdp=True
          )
          results = run_integration_test(pretrain_config)
          losses = results.get_losses()
          print(losses)
  
          # Convert checkpoint to sliced checkpoint for easy loading
          CheckpointFormatConverter.sharded_to_sliced_checkpoint(
              "checkpoint.torch", "checkpoint_sliced.torch"
          )
          checkpoint_path = os.path.join(
              pretrain_dir, "checkpoint_sliced.torch"
          )
  
          # Create a directory to contain the extracted features
          with in_temporary_directory() as extract_dir:
              # Extract head features
              extract_config_head = (
                  self._create_extract_features_config_head(
                      checkpoint_path=checkpoint_path, with_fsdp=True
                  )
              )
              extract_config_head.EXTRACT_FEATURES.OUTPUT_DIR = (
                  extract_dir
              )
              run_integration_test(
                  extract_config_head, engine_name="extract_features"
              )
  
              # Extract trunk features
              extract_config_trunk = (
                  self._create_extract_features_config_trunk(
                      checkpoint_path=checkpoint_path, with_fsdp=True
                  )
              )
              extract_config_trunk.EXTRACT_FEATURES.OUTPUT_DIR = (
                  extract_dir
              )
              run_integration_test(
                  extract_config_trunk, engine_name="extract_features"
              )
  
              # Verify that we can merge the heads features back
              train_feat = ExtractedFeaturesLoader.load_features(
                  extract_dir, "train", "heads", flatten_features=True
              )
              self.assertEqual(
                  train_feat["features"].shape, torch.Size([200, 128])
              )
              self.assertEqual(
                  train_feat["targets"].shape, torch.Size([200, 1])
              )
              self.assertEqual(
                  train_feat["inds"].shape, torch.Size([200])
              )
  
              # Verify that we can merge the trunk features back
              train_feat = ExtractedFeaturesLoader.load_features(
                  extract_dir, "train", "res5", flatten_features=True
              )
              self.assertEqual(
                  train_feat["features"].shape,
                  torch.Size([200, 3024 * 2 * 2]),
              )
              self.assertEqual(
                  train_feat["targets"].shape, torch.Size([200, 1])
              )
              self.assertEqual(
                  train_feat["inds"].shape, torch.Size([200])
              )
  
              # Run KNN on the res5 layer
              extract_config_trunk.NEAREST_NEIGHBOR.FEATURES.PATH = (
                  extract_dir
              )
              top_1_ref, top_5_ref, total_ref = run_knn_at_layer(
                  extract_config_trunk, layer_name="res5"
              )
              (
                  top_1_opt,
                  top_5_opt,
                  total_opt,
              ) = run_knn_at_layer_low_memory(
                  extract_config_trunk, layer_name="res5"
              )
              self.assertEqual(total_ref, total_opt)
              # TODO - investigate: both KNN implementation have a bit of randomness
              #  in their accuracies, so the asserts are inequalities.
              self.assertLessEqual(top_1_ref, 30.0)
              self.assertLessEqual(top_1_opt, 30.0)
              self.assertGreaterEqual(top_1_ref, 29.0)
              self.assertGreaterEqual(top_1_opt, 29.0)
              # self.assertEqual(top_1_ref, top_1_opt)
              # self.assertEqual(top_5_ref, top_5_opt)
  
              # Run KNN on the head layer
              extract_config_head.NEAREST_NEIGHBOR.FEATURES.PATH = (
                  extract_dir
              )
              top_1_ref, top_5_ref, total_ref = run_knn_at_layer(
                  extract_config_head, layer_name="heads"
              )
              (
                  top_1_opt,
                  top_5_opt,
                  total_opt,
              ) = run_knn_at_layer_low_memory(
                  extract_config_head, layer_name="heads"
              )
              self.assertEqual(total_ref, total_opt)
              # TODO - investigate: both KNN implementation have a bit of randomness
              #  in their accuracies, so the asserts are inequalities.
              self.assertLessEqual(top_1_ref, 35.0)
              self.assertLessEqual(top_1_opt, 35.0)
              self.assertGreaterEqual(top_1_ref, 33.0)
              self.assertGreaterEqual(top_1_opt, 33.0)
  
file_path: vissl/tests/test_knn.py
imports: import os
  import unittest
  import torch
  from vissl.utils.checkpoint import CheckpointFormatConverter
  from vissl.utils.extract_features_utils import ExtractedFeaturesLoader
  from vissl.utils.hydra_config import compose_hydra_configuration, convert_to_attrdict
  from vissl.utils.knn_utils import run_knn_at_layer, run_knn_at_layer_low_memory
  from vissl.utils.test_utils import gpu_test, in_temporary_directory, run_integration_test
label_desc: Use image node_id 6 for calling the TestKNNBenchmark obj's underlying member method code with example usage: obj.test_knn_fsdp() without return types