Architecture

• main.py
• Model
  • cae_model.py
  • vcae_model.py
  • refiner.py
  • hybrid.py: The refiner that combine transformer and mamba. (Mostly generated by LLM, not tested thoroughly.)
  • diffusion_model.py: diffusion model as refiner
    • To use diffusion as decoder, modify model_type.
    • To make training more efficient, reduce the step of inference (and maybe the dimension), since current training method requires the whole inference process during training.
• Auxiliary
  • dataset.py: Normalization method can be changed.
  • plot.py: heat-map and channel plot
  • utils.py
    • Define the type of refiner.
    • Calculation of scores.
    • Abstraction for getting compressed code/decoding.
    • Fix seed.

Usage

1. Install the required package pip install -r requirement.txt
2. Install the package for mamba properly, if there is system issue,
   • try reinstall pytroch properly. https://pytorch.org/get-started/locally/
   • try to install from the original implementation https://github.com/state-spaces/mamba/tree/main, or the package provided by https://github.com/khhungg/MECG-E/tree/main (this worked on iis server when I tried).
3. In main.py
   • Choose the refiner to use (hyperparameter of mamba and hybrid not optimized), and which subset of the data to train.
   • Modify the path of model and the directory to store result.
4. Run python3 main.py name_of_testing