This is a repository for building Wireless Signal Recognition (WSR) foundation models based on PyTorch, integrating various signal datasets, models, and model experiment pipelines. We try to provide a user-transparent and easy-to-use benchmarking interface.
Note: This version is currently only used as an implementation reference for TS-DDAE [ICLR 2026]. Updates to this repository will be made gradually in the future.
- PyTorch >= 2.2.0
- sklearn
- scipy
- pandas
- cv2
- Optuna (if you want hyperparameters optimization)
- matplotlib, seaborn (if you want visualization)
python main.py -m model_name -d dataset_name -t task_name -g 0 --load_from_pretrainedoptional arguments:
--model -m name of models
--task -t name of task
--dataset -d name of datasets
--gpu -g controls which gpu you will use. If you do not have gpu, set -g -1.
--load_from_pretrained will load the model from a default checkpoint.
--compile will use the torch.compile optimization. (still under development)
Note: some usage examples are provided in "/example/".
We provide two ways of modifying the hyperparameters of the model
- By specifying parameters in the experiment function in
main.py. - By modifying the
config.yamlfile in each model folder in the foundwsr/models.
Note: the first method has higher priority than the second.
If you want to add new pipelines, you need to
-
Inherit the BasePipe from
foundwsr.pipeline.base_pipe. -
Register the pipeline name using
register_pipeinfoundwsr.pipeline.__init__. -
Add the pipeline name and file path to the
SUPPORTED_PIPESvariable infoundwsr.pipeline.__init__. -
Write the pipeline's operation logic.
-
Add the task name and pipeline name to
supported.yamlin the foundwsr folder.
The pipeline calling logic has been encapsulated in the experiment.py file in the foundwsr folder.
Similar to the pipeline support, if you want to add new datasets, you need to
-
Inherit the BaseDataset from
foundwsr.dataset.base_dataset. -
Register the dataset name using
register_datasetinfoundwsr.dataset.__init__. -
Add the dataset name and file path to the
SUPPORTED_DATASETSvariable infoundwsr.dataset.__init__. -
Write the dataset's operation logic.
Note: we recommend you to write the @classmethod create if you want to enjoy the build_dataset from foundwsr.dataset.__init__. In create, users can define the initial data loading and data partitioning. For task-specific dataset like MaskedDataset, you may not write the create.
Similar to the dataset support, if you want to add new models, you need to
-
Create a folder in the
foundwsr.models. -
Inherit the BaseModel from
foundwsr.models.base_model. -
Register the model name using
register_modelinfoundwsr.models.__init__. -
Add the model name and file path to the
SUPPORTED_MODELSvariable infoundwsr.models.__init__. -
Write the model's operation logic.
-
Add
config.yamlto specify the hyperparameters used in this model.
Note: we recommend you to write the @classmethod build_model_from_args if you want to enjoy the build_model from foundwsr.model.__init__. In build_model_from_args, users can define the initial input hyperparameters.
This repository is still under development. If you have any questions, suggestions, or would like to contribute, please open an issue, submit a pull request, or contact me directly at yaoqiliu@bupt.edu.cn