update

Author: asif-hanif

DATASETS.md

@@ -63,7 +63,7 @@ Each CSV file includes the following columns:
 <hr><hr>
 <br><br>
 
-We have uploaded all datasets on [Huggingface Datasets](https://huggingface.co/docs/datasets/en/index). Following are the python commands to download datasets. Make sure to provide valid destination dataset path ending with 'Audio-Datasets' folder and install `huggingface_hub` package. We have also provided a [Jupyter Notebook](/utils/DownloadAudioDatasets.ipynb) to download all datasets in one go. It might take some time to download all datasets, so we recommend running the notebook on a cloud instance or a machine with good internet speed.<br><br>
+We have uploaded all datasets on [Huggingface Datasets](https://huggingface.co/docs/datasets/en/index). Following are the python commands to download datasets. Make sure to provide valid destination dataset path ending with 'Audio-Datasets' folder and install `huggingface_hub` package. We have also provided a [Jupyter Notebook](/media/DownloadAudioDatasets.ipynb) to download all datasets in one go. It might take some time to download all datasets, so we recommend running the notebook on a cloud instance or a machine with good internet speed.<br><br>
 `pip install huggingface-hub==0.25.1`
 
 <br>

README.md

@@ -94,7 +94,7 @@ wget https://zenodo.org/records/8387083/files/base.pth
 
 ## Datasets :page_with_curl:
 
-We have performed experiments on 11 audio classification datasets.  Instructions for downloading/processing datasets used by our method have been provided in the [DATASETS.md](DATASETS.md). All of the datasets have been uploaded on HuggingFace Datasets Hub :hugs: for easy access. We have also provided a [Jupyter Notebook](/utils/DownloadAudioDatasets.ipynb) to download all datasets in one go. It might take some time to download all datasets, so we recommend running the notebook on a cloud instance or a machine with good internet speed.
+We have performed experiments on 11 audio classification datasets.  Instructions for downloading/processing datasets used by our method have been provided in the [DATASETS.md](DATASETS.md). All of the datasets have been uploaded on HuggingFace Datasets Hub :hugs: for easy access. We have also provided a [Jupyter Notebook](/media/DownloadAudioDatasets.ipynb) to download all datasets in one go. It might take some time to download all datasets, so we recommend running the notebook on a cloud instance or a machine with good internet speed.
 
 | Dataset | Type | Classes | Size | Link |
 |:-- |:-- |:--: |--: |:-- |

utils/DownloadAudioDatasets.ipynb media/DownloadAudioDatasets.ipynb

@@ -41,14 +41,28 @@ def get_dataloaders(args):
 
 
 def save_model(args, model, save_model_path):
-        print(f"Saving Context Weights for Method: '{args.model_name.upper()}'\n")
-        if args.model_name in ['coop', 'cocoop', 'palm']:
-            torch.save(model.prompt_learner.state_dict(), save_model_path)
-        else:
-            raise ValueError(f"Model '{args.model_name}' is not supported. Choose from: [{', '.join(METHODS)}]")
-        
+    print(f"Saving Context Weights for Method: '{args.model_name.upper()}'\n")
+    if args.model_name in ['coop', 'cocoop', 'palm']:
+        checkpoint = {'prompt_learner': model.prompt_learner.state_dict()}
+        checkpoint['pengi_bn0_buffer'] = {'running_mean': model.audio_encoder.base.htsat.bn0.running_mean.clone(),
+                                            'running_var': model.audio_encoder.base.htsat.bn0.running_var.clone(),
+                                            'num_batches_tracked': model.audio_encoder.base.htsat.bn0.num_batches_tracked.clone()}
+
+        torch.save(checkpoint, save_model_path)
+    else:
+        raise ValueError(f"Model '{args.model_name}' is not supported. Choose from: [{', '.join(METHODS)}]")
+
+
+    
 def load_model(args, model):
-        raise NotImplementedError("\n\nLoading model is not implemented yet.\n\n")
+        load_model_path = get_load_model_path(args)
+        checkpoint = torch.load(load_model_path)
+        model.prompt_learner.load_state_dict(checkpoint['prompt_learner'])
+        model.audio_encoder.base.htsat.bn0.running_mean.copy_(checkpoint['pengi_bn0_buffer']['running_mean'])
+        model.audio_encoder.base.htsat.bn0.running_var.copy_(checkpoint['pengi_bn0_buffer']['running_var'])
+        model.audio_encoder.base.htsat.bn0.num_batches_tracked.copy_(checkpoint['pengi_bn0_buffer']['num_batches_tracked'])
+        # raise NotImplementedError("\n\nLoading model is not implemented yet.\n\n")
+
 
 
 def get_save_model_path(args):
@@ -105,8 +119,10 @@ def get_args():
         raise ValueError(f"\n\nDirectory '{args.dataset_root}' does not exist. Specify the correct path to the dataset.\n\n")
     if args.save_model and not os.path.exists(args.save_model_path):
         raise ValueError(f"\n\nDirectory '{args.save_model_path}' does not exist. Create or specify the correct the directory to save the trained model.\n\n")
-    if args.eval_only and not os.path.exists(args.load_model_path):
-        raise ValueError(f"\n\nEvaluation Mode: Model file '{args.load_model_path}' does not exist. Specify the correct path to the model file.\n\n")
+    if args.eval_only:
+        load_model_path = get_load_model_path(args)
+        if not os.path.exists(load_model_path): raise ValueError(f"\n\nEvaluation Mode: Model file '{load_model_path}' does not exist. Specify the correct path to the model file.\n\n")
+    
     if args.model_name == 'zeroshot': args.eval_only = True
 
     return args