Braindecode 1.3.1.dev177756420__tar.gz → 1.3.1.dev178641268__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (161) hide show
  1. {braindecode-1.3.1.dev177756420/braindecode.egg-info → braindecode-1.3.1.dev178641268}/PKG-INFO +1 -1
  2. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/models/__init__.py +0 -2
  3. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/models/base.py +1 -10
  4. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/models/summary.csv +0 -1
  5. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/models/util.py +0 -1
  6. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/modules/__init__.py +0 -2
  7. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/modules/attention.py +0 -134
  8. braindecode-1.3.1.dev178641268/braindecode/version.py +1 -0
  9. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268/braindecode.egg-info}/PKG-INFO +1 -1
  10. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode.egg-info/SOURCES.txt +0 -1
  11. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/docs/api.rst +10 -11
  12. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/docs/whats_new.rst +0 -1
  13. braindecode-1.3.1.dev177756420/braindecode/models/cbramod.py +0 -350
  14. braindecode-1.3.1.dev177756420/braindecode/version.py +0 -1
  15. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/LICENSE.txt +0 -0
  16. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/MANIFEST.in +0 -0
  17. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/NOTICE.txt +0 -0
  18. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/README.rst +0 -0
  19. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/__init__.py +0 -0
  20. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/augmentation/__init__.py +0 -0
  21. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/augmentation/base.py +0 -0
  22. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/augmentation/functional.py +0 -0
  23. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/augmentation/transforms.py +0 -0
  24. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/classifier.py +0 -0
  25. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/datasets/__init__.py +0 -0
  26. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/datasets/base.py +0 -0
  27. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/datasets/bbci.py +0 -0
  28. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/datasets/bcicomp.py +0 -0
  29. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/datasets/bids/__init__.py +0 -0
  30. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/datasets/bids/datasets.py +0 -0
  31. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/datasets/bids/format.py +0 -0
  32. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/datasets/bids/hub.py +0 -0
  33. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/datasets/bids/hub_format.py +0 -0
  34. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/datasets/bids/hub_io.py +0 -0
  35. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/datasets/bids/hub_validation.py +0 -0
  36. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/datasets/bids/iterable.py +0 -0
  37. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/datasets/chb_mit.py +0 -0
  38. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/datasets/mne.py +0 -0
  39. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/datasets/moabb.py +0 -0
  40. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/datasets/nmt.py +0 -0
  41. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/datasets/registry.py +0 -0
  42. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/datasets/siena.py +0 -0
  43. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/datasets/sleep_physio_challe_18.py +0 -0
  44. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/datasets/sleep_physionet.py +0 -0
  45. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/datasets/tuh.py +0 -0
  46. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/datasets/utils.py +0 -0
  47. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/datasets/xy.py +0 -0
  48. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/datautil/__init__.py +0 -0
  49. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/datautil/channel_utils.py +0 -0
  50. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/datautil/hub_formats.py +0 -0
  51. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/datautil/serialization.py +0 -0
  52. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/datautil/util.py +0 -0
  53. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/eegneuralnet.py +0 -0
  54. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/functional/__init__.py +0 -0
  55. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/functional/functions.py +0 -0
  56. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/functional/initialization.py +0 -0
  57. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/models/atcnet.py +0 -0
  58. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/models/attentionbasenet.py +0 -0
  59. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/models/attn_sleep.py +0 -0
  60. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/models/bendr.py +0 -0
  61. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/models/biot.py +0 -0
  62. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/models/brainmodule.py +0 -0
  63. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/models/config.py +0 -0
  64. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/models/contrawr.py +0 -0
  65. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/models/ctnet.py +0 -0
  66. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/models/deep4.py +0 -0
  67. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/models/deepsleepnet.py +0 -0
  68. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/models/eegconformer.py +0 -0
  69. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/models/eeginception_erp.py +0 -0
  70. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/models/eeginception_mi.py +0 -0
  71. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/models/eegitnet.py +0 -0
  72. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/models/eegminer.py +0 -0
  73. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/models/eegnet.py +0 -0
  74. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/models/eegnex.py +0 -0
  75. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/models/eegpt.py +0 -0
  76. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/models/eegsimpleconv.py +0 -0
  77. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/models/eegsym.py +0 -0
  78. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/models/eegtcnet.py +0 -0
  79. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/models/fbcnet.py +0 -0
  80. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/models/fblightconvnet.py +0 -0
  81. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/models/fbmsnet.py +0 -0
  82. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/models/hybrid.py +0 -0
  83. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/models/ifnet.py +0 -0
  84. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/models/labram.py +0 -0
  85. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/models/luna.py +0 -0
  86. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/models/medformer.py +0 -0
  87. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/models/msvtnet.py +0 -0
  88. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/models/patchedtransformer.py +0 -0
  89. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/models/reve.py +0 -0
  90. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/models/sccnet.py +0 -0
  91. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/models/shallow_fbcsp.py +0 -0
  92. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/models/signal_jepa.py +0 -0
  93. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/models/sinc_shallow.py +0 -0
  94. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/models/sleep_stager_blanco_2020.py +0 -0
  95. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/models/sleep_stager_chambon_2018.py +0 -0
  96. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/models/sparcnet.py +0 -0
  97. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/models/sstdpn.py +0 -0
  98. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/models/syncnet.py +0 -0
  99. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/models/tcn.py +0 -0
  100. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/models/tidnet.py +0 -0
  101. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/models/tsinception.py +0 -0
  102. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/models/usleep.py +0 -0
  103. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/modules/activation.py +0 -0
  104. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/modules/blocks.py +0 -0
  105. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/modules/convolution.py +0 -0
  106. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/modules/filter.py +0 -0
  107. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/modules/layers.py +0 -0
  108. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/modules/linear.py +0 -0
  109. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/modules/parametrization.py +0 -0
  110. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/modules/stats.py +0 -0
  111. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/modules/util.py +0 -0
  112. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/modules/wrapper.py +0 -0
  113. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/preprocessing/__init__.py +0 -0
  114. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/preprocessing/eegprep_preprocess.py +0 -0
  115. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/preprocessing/mne_preprocess.py +0 -0
  116. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/preprocessing/preprocess.py +0 -0
  117. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/preprocessing/util.py +0 -0
  118. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/preprocessing/windowers.py +0 -0
  119. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/regressor.py +0 -0
  120. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/samplers/__init__.py +0 -0
  121. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/samplers/base.py +0 -0
  122. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/samplers/ssl.py +0 -0
  123. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/training/__init__.py +0 -0
  124. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/training/callbacks.py +0 -0
  125. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/training/losses.py +0 -0
  126. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/training/scoring.py +0 -0
  127. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/util.py +0 -0
  128. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/visualization/__init__.py +0 -0
  129. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/visualization/confusion_matrices.py +0 -0
  130. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode/visualization/gradients.py +0 -0
  131. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode.egg-info/dependency_links.txt +0 -0
  132. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode.egg-info/requires.txt +0 -0
  133. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/braindecode.egg-info/top_level.txt +0 -0
  134. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/docs/Makefile +0 -0
  135. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/docs/_templates/autosummary/class.rst +0 -0
  136. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/docs/_templates/autosummary/class_in_subdir.rst +0 -0
  137. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/docs/_templates/autosummary/function.rst +0 -0
  138. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/docs/_templates/autosummary/function_in_subdir.rst +0 -0
  139. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/docs/cite.rst +0 -0
  140. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/docs/conf.py +0 -0
  141. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/docs/help.rst +0 -0
  142. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/docs/index.rst +0 -0
  143. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/docs/install/install.rst +0 -0
  144. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/docs/install/install_pip.rst +0 -0
  145. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/docs/install/install_source.rst +0 -0
  146. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/docs/models/categorization/attention.rst +0 -0
  147. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/docs/models/categorization/channel.rst +0 -0
  148. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/docs/models/categorization/convolution.rst +0 -0
  149. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/docs/models/categorization/filterbank.rst +0 -0
  150. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/docs/models/categorization/gnn.rst +0 -0
  151. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/docs/models/categorization/interpretable.rst +0 -0
  152. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/docs/models/categorization/lbm.rst +0 -0
  153. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/docs/models/categorization/recurrent.rst +0 -0
  154. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/docs/models/categorization/spd.rst +0 -0
  155. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/docs/models/models.rst +0 -0
  156. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/docs/models/models_categorization.rst +0 -0
  157. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/docs/models/models_table.rst +0 -0
  158. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/docs/models/models_visualization.rst +0 -0
  159. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/docs/sg_execution_times.rst +0 -0
  160. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/pyproject.toml +0 -0
  161. {braindecode-1.3.1.dev177756420 → braindecode-1.3.1.dev178641268}/setup.cfg +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: braindecode
3
- Version: 1.3.1.dev177756420
3
+ Version: 1.3.1.dev178641268
4
4
  Summary: Deep learning software to decode EEG, ECG or MEG signals
5
5
  Author-email: Robin Tibor Schirrmeister <robintibor@gmail.com>, Bruno Aristimunha Pinto <b.aristimunha@gmail.com>, Alexandre Gramfort <agramfort@meta.com>
6
6
  Maintainer-email: Alexandre Gramfort <agramfort@meta.com>, Bruno Aristimunha Pinto <b.aristimunha@gmail.com>, Robin Tibor Schirrmeister <robintibor@gmail.com>
@@ -7,7 +7,6 @@ from .base import EEGModuleMixin
7
7
  from .bendr import BENDR
8
8
  from .biot import BIOT
9
9
  from .brainmodule import BrainModule
10
- from .cbramod import CBraMod
11
10
  from .contrawr import ContraWR
12
11
  from .ctnet import CTNet
13
12
  from .deep4 import Deep4Net
@@ -69,7 +68,6 @@ __all__ = [
69
68
  "EEGModuleMixin",
70
69
  "BIOT",
71
70
  "BENDR",
72
- "CBraMod",
73
71
  "ContraWR",
74
72
  "CTNet",
75
73
  "Deep4Net",
@@ -397,18 +397,9 @@ class EEGModuleMixin(_BaseHubMixin, metaclass=NumpyDocstringInheritanceInitMeta)
397
397
  torchinfo.ModelStatistics
398
398
  ModelStatistics generated by torchinfo.summary.
399
399
  """
400
- try:
401
- n_chans = self.n_chans
402
- except ValueError:
403
- n_chans = 1
404
- try:
405
- n_times = self.n_times
406
- except ValueError:
407
- n_times = 200
408
-
409
400
  return summary(
410
401
  self,
411
- input_size=(1, n_chans, n_times),
402
+ input_size=(1, self.n_chans, self.n_times),
412
403
  col_names=col_names,
413
404
  row_settings=row_settings,
414
405
  verbose=0,
@@ -3,7 +3,6 @@ ATCNet,General,Classification,250,"n_chans, n_outputs, n_times",113732,"ATCNet(n
3
3
  AttentionBaseNet,Motor Imagery,Classification,250,"n_chans, n_outputs, n_times",3692,"AttentionBaseNet(n_chans=22, n_outputs=4, n_times=1000)","Convolution,Attention/Transformer"
4
4
  BDTCN,Normal Abnormal,Classification,100,"n_chans, n_outputs, n_times",456502,"BDTCN(n_chans=21, n_outputs=2, n_times=6000, n_blocks=5, n_filters=55, kernel_size=16)","Convolution,Recurrent"
5
5
  BIOT,"Sleep Staging, Epilepsy",Classification,200,"n_chans, n_outputs",3183879,"BIOT(n_chans=2, n_outputs=5, n_times=6000)","Foundation Model"
6
- CBraMod,General,"Classification,Embedding",200,"n_outputs",4924000,"CBraMod(n_outputs=2)","Foundation Model"
7
6
  ContraWR,Sleep Staging,"Classification, Embedding",125,"n_chans, n_outputs, sfreq",1160165,"ContraWR(n_chans=2, n_outputs=5, n_times=3750, emb_size=256, sfreq=125)",Convolution
8
7
  CTNet,Motor Imagery,Classification,250,"n_chans, n_outputs, n_times",26900,"CTNet(n_chans=22, n_outputs=4, n_times=1000, n_filters_time=8, kernel_size=16, num_heads=2, embed_dim=16)","Convolution,Attention/Transformer"
9
8
  Deep4Net,General,Classification,250,"n_chans, n_outputs, n_times",282879,"Deep4Net(n_chans=22, n_outputs=4, n_times=1000)","Convolution"
@@ -132,7 +132,6 @@ models_mandatory_parameters: list[
132
132
  "chs_info": [{"ch_name": f"E{i + 1}", "kind": "eeg"} for i in range(19)],
133
133
  },
134
134
  ),
135
- ("CBraMod", ["n_outputs"], None),
136
135
  ]
137
136
 
138
137
  ################################################################
@@ -7,7 +7,6 @@ from .attention import (
7
7
  GCT,
8
8
  SRM,
9
9
  CATLite,
10
- CrissCrossTransformerEncoderLayer,
11
10
  EncNet,
12
11
  GatherExcite,
13
12
  GSoP,
@@ -60,7 +59,6 @@ __all__ = [
60
59
  "GSoP",
61
60
  "MultiHeadAttention",
62
61
  "SqueezeAndExcitation",
63
- "CrissCrossTransformerEncoderLayer",
64
62
  "MLP",
65
63
  "FeedForwardBlock",
66
64
  "InceptionBlock",
@@ -881,137 +881,3 @@ class MultiHeadAttention(nn.Module):
881
881
  out = self.rearrange_unstack(out)
882
882
  out = self.projection(out)
883
883
  return out
884
-
885
-
886
- class CrissCrossTransformerEncoderLayer(nn.Module):
887
- """
888
- **Criss-Cross Transformer Encoder Layer** implementing the criss-cross attention mechanism
889
- proposed in Wang et al. (2025) [cbramod]_.
890
-
891
- This layer models spatial and temporal dependencies in EEG signals through parallel
892
- **Spatial Attention** (S-Attention) and **Temporal Attention** (T-Attention). Unlike
893
- standard Transformer layers that model all spatial-temporal dependencies together,
894
- this layer separates these heterogeneous dependencies:
895
-
896
- - **S-Attention**: Captures dependencies between EEG channels within a fixed time interval.
897
- For each temporal stripe across all channels, it applies self-attention independently,
898
- learning channel interactions while keeping time steps separate.
899
-
900
- - **T-Attention**: Captures dependencies between temporal patches within a fixed channel.
901
- For each channel across all time intervals, it applies self-attention independently,
902
- learning temporal dynamics while keeping channels separate.
903
-
904
- The attention heads are split equally between S-Attention and T-Attention, and their
905
- outputs are concatenated. This criss-cross design effectively captures EEG's unique
906
- structural characteristics—EEG signals exhibit heterogeneous spatial and temporal
907
- dependencies that differ from images (which have only spatial dependencies).
908
-
909
- Input shape: ``(batch, channels, n_patches, patch_size)``
910
- Output shape: ``(batch, channels, n_patches, patch_size)``
911
-
912
- .. figure:: https://braindecode.org/dev/_static/model/criss-cross-attention.png
913
- :align: center
914
- :alt: Criss-Cross Attention Mechanism
915
-
916
- References
917
- ----------
918
- .. [cbramod] Wang, J., Zhao, S., Luo, Z., Zhou, Y., Jiang, H., Li, S., Li, T., & Pan, G. (2025).
919
- CBraMod: A Criss-Cross Brain Foundation Model for EEG Decoding.
920
- In The Thirteenth International Conference on Learning Representations (ICLR 2025).
921
- https://arxiv.org/abs/2412.07236
922
- """
923
-
924
- __constants__ = ["norm_first"]
925
-
926
- def __init__(
927
- self,
928
- d_model: int,
929
- nhead: int,
930
- dim_feedforward: int = 2048,
931
- dropout: float = 0.1,
932
- activation: type[nn.Module] = nn.GELU,
933
- layer_norm_eps: float = 1e-5,
934
- batch_first: bool = False,
935
- norm_first: bool = False,
936
- bias: bool = True,
937
- ) -> None:
938
- super().__init__()
939
- self.self_attn_s = nn.MultiheadAttention(
940
- d_model // 2,
941
- nhead // 2,
942
- dropout=dropout,
943
- bias=bias,
944
- batch_first=batch_first,
945
- )
946
- self.self_attn_t = nn.MultiheadAttention(
947
- d_model // 2,
948
- nhead // 2,
949
- dropout=dropout,
950
- bias=bias,
951
- batch_first=batch_first,
952
- )
953
-
954
- # Implementation of Feedforward model
955
- self.linear1 = nn.Linear(d_model, dim_feedforward, bias=bias)
956
- self.dropout = nn.Dropout(dropout)
957
- self.linear2 = nn.Linear(dim_feedforward, d_model, bias=bias)
958
-
959
- self.norm_first = norm_first
960
- self.norm1 = nn.LayerNorm(d_model, eps=layer_norm_eps)
961
- self.norm2 = nn.LayerNorm(d_model, eps=layer_norm_eps)
962
- self.dropout1 = nn.Dropout(dropout)
963
- self.dropout2 = nn.Dropout(dropout)
964
- self.activation = activation()
965
-
966
- def forward(
967
- self,
968
- src: Tensor,
969
- src_mask: Optional[Tensor] = None,
970
- src_key_padding_mask: Optional[Tensor] = None,
971
- is_causal: bool = False,
972
- ) -> Tensor:
973
- x = src
974
- x = x + self._sa_block(
975
- self.norm1(x), src_mask, src_key_padding_mask, is_causal=is_causal
976
- )
977
- x = x + self._ff_block(self.norm2(x))
978
- return x
979
-
980
- # self-attention block
981
- def _sa_block(
982
- self,
983
- x: Tensor,
984
- attn_mask: Optional[Tensor],
985
- key_padding_mask: Optional[Tensor],
986
- is_causal: bool = False,
987
- ) -> Tensor:
988
- bz, ch_num, patch_num, patch_size = x.shape
989
- xs = x[:, :, :, : patch_size // 2]
990
- xt = x[:, :, :, patch_size // 2 :]
991
- xs = rearrange(xs, "b c n p2 -> (b n) c p2")
992
- xt = rearrange(xt, "b c n p2 -> (b c) n p2")
993
- xs = self.self_attn_s(
994
- xs,
995
- xs,
996
- xs,
997
- attn_mask=attn_mask,
998
- key_padding_mask=key_padding_mask,
999
- need_weights=False,
1000
- )[0]
1001
- xs = rearrange(xs, "(b n) c p2 -> b c n p2", b=bz, n=patch_num)
1002
- xt = self.self_attn_t(
1003
- xt,
1004
- xt,
1005
- xt,
1006
- attn_mask=attn_mask,
1007
- key_padding_mask=key_padding_mask,
1008
- need_weights=False,
1009
- )[0]
1010
- xt = rearrange(xt, "(b c) n p2 -> b c n p2", b=bz, c=ch_num)
1011
- x = torch.concat((xs, xt), dim=3)
1012
- return self.dropout1(x)
1013
-
1014
- # feed forward block
1015
- def _ff_block(self, x: Tensor) -> Tensor:
1016
- x = self.linear2(self.dropout(self.activation(self.linear1(x))))
1017
- return self.dropout2(x)
@@ -0,0 +1 @@
1
+ __version__ = "1.3.1.dev178641268"
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: braindecode
3
- Version: 1.3.1.dev177756420
3
+ Version: 1.3.1.dev178641268
4
4
  Summary: Deep learning software to decode EEG, ECG or MEG signals
5
5
  Author-email: Robin Tibor Schirrmeister <robintibor@gmail.com>, Bruno Aristimunha Pinto <b.aristimunha@gmail.com>, Alexandre Gramfort <agramfort@meta.com>
6
6
  Maintainer-email: Alexandre Gramfort <agramfort@meta.com>, Bruno Aristimunha Pinto <b.aristimunha@gmail.com>, Robin Tibor Schirrmeister <robintibor@gmail.com>
@@ -58,7 +58,6 @@ braindecode/models/base.py
58
58
  braindecode/models/bendr.py
59
59
  braindecode/models/biot.py
60
60
  braindecode/models/brainmodule.py
61
- braindecode/models/cbramod.py
62
61
  braindecode/models/config.py
63
62
  braindecode/models/contrawr.py
64
63
  braindecode/models/ctnet.py
@@ -62,16 +62,16 @@ a common interface for all EEG models and derivate variables names if necessary.
62
62
  Some models have pre-trained weights available on the Hugging Face BrainDecode
63
63
  organization:
64
64
 
65
- .. currentmodule:: braindecode.models
66
-
67
- - :class:`BIOT` - Foundation model with pre-trained weights
68
- - :class:`CBraMod` - Criss-Cross Transformer model with pre-trained weights
69
- - :class:`Labram` - Large Brain Model with pre-trained weights
70
- - :class:`REVE` - EEG foundation model with pre-trained weights
71
- - :class:`LUNA` - Universal EEG embedding model with pre-trained weights
72
- - :class:`BENDR` - Foundation model with pre-trained weights
73
- - :class:`SignalJEPA` - Self-supervised learning model with pre-trained weights
74
- - :class:`EEGPT` - Pretrained transformer for universal EEG
65
+ - :class:`~braindecode.models.BIOT` - Foundation model with pre-trained weights
66
+ - :class:`~braindecode.models.Labram` - Large Brain Model with pre-trained weights
67
+ - :class:`~braindecode.models.REVE` - EEG foundation model with pre-trained weights
68
+ - :class:`~braindecode.models.LUNA` - Universal EEG embedding model with pre-trained
69
+ weights
70
+ - :class:`~braindecode.models.BENDR` - Foundation model with pre-trained weights
71
+ - :class:`~braindecode.models.SignalJEPA` - Self-supervised learning model with
72
+ pre-trained weights
73
+ - :class:`~braindecode.models.EEGPT` - Pretrained transformer for universal EEG
74
+ representation
75
75
 
76
76
  **Example - Loading a pre-trained model:**
77
77
 
@@ -132,7 +132,6 @@ a common interface for all EEG models and derivate variables names if necessary.
132
132
  BENDR
133
133
  BIOT
134
134
  BrainModule
135
- CBraMod
136
135
  ContraWR
137
136
  CTNet
138
137
  Deep4Net
@@ -33,7 +33,6 @@ Enhancements
33
33
  - Adding :class:`braindecode.models.EEGPT` with HuggingFace Hub integration and
34
34
  pretrained weights available at `braindecode/eegpt-pretrained <https://huggingface.co/braindecode/eegpt-pretrained>`_
35
35
  (:gh:`908` by `Young Truong`_ and `Kuntal Kokate`_)
36
- - Add :class:`braindecode.models.CBraMod` model (:gh:`914` by `Pierre Guetschel`_)
37
36
 
38
37
  API changes
39
38
  ============
@@ -1,350 +0,0 @@
1
- # Authors: Pierre Guetschel
2
- #
3
- # Code adapted from https://github.com/wjq-learning/CBraMod
4
- #
5
- # License: BSD (3-clause)
6
-
7
-
8
- import copy
9
- import logging
10
- from typing import Optional, Sequence
11
-
12
- import torch
13
- from einops import rearrange
14
- from einops.layers.torch import Rearrange
15
- from torch import Tensor, nn
16
-
17
- from braindecode.models.base import EEGModuleMixin
18
- from braindecode.modules import CrissCrossTransformerEncoderLayer
19
-
20
- logger = logging.getLogger(__name__)
21
-
22
-
23
- class CBraMod(EEGModuleMixin, nn.Module):
24
- r"""
25
- **C**\ riss-\ **C**\ ross **Bra**\ in **Mod**\ el for EEG Decoding from Wang et al. (2025) [cbramod]_.
26
-
27
- :bdg-danger:`Foundation Model` :bdg-info:`Attention/Transformer`
28
-
29
- .. figure:: https://raw.githubusercontent.com/wjq-learning/CBraMod/refs/heads/main/figure/model.png
30
- :align: center
31
- :alt: CBraMod pre-training overview
32
- :width: 1000px
33
-
34
- CBraMod is a foundation model for EEG decoding that leverages a novel criss-cross transformer
35
- architecture to effectively model the unique spatial and temporal characteristics of EEG signals.
36
- Pre-trained on the Temple University Hospital EEG Corpus (TUEG)—the largest public EEG corpus—
37
- using masked EEG patch reconstruction, CBraMod achieves state-of-the-art performance across
38
- diverse downstream BCI and clinical applications.
39
-
40
- .. rubric:: Key Innovation: Criss-Cross Attention
41
-
42
- Unlike existing EEG foundation models that use full attention to model all spatial and temporal
43
- dependencies together, CBraMod separates spatial and temporal dependencies through a
44
- **criss-cross transformer** architecture:
45
-
46
- - **Spatial Attention**: Models dependencies between channels while keeping patches separate
47
- - **Temporal Attention**: Models dependencies between temporal patches while keeping channels separate
48
-
49
- This design is inspired by criss-cross strategies from computer vision and effectively
50
- leverages the inherent structural characteristics of EEG signals. The criss-cross approach
51
- reduces computational complexity (FLOPs reduced by ~32% compared to full attention) while
52
- improving performance and enabling faster convergence.
53
-
54
- .. rubric:: Asymmetric Conditional Positional Encoding (ACPE)
55
-
56
- Rather than using fixed positional embeddings, CBraMod employs **Asymmetric Conditional
57
- Positional Encoding** that dynamically generates positional embeddings using a convolutional
58
- network. This enables the model to:
59
-
60
- - Capture relative positional information adaptively
61
- - Handle diverse EEG channel formats (different channel counts and reference schemes)
62
- - Generalize to arbitrary downstream EEG formats without retraining
63
- - Support various reference schemes (earlobe, average, REST, bipolar)
64
-
65
- .. rubric:: Pretraining Highlights
66
-
67
- - **Pretraining Dataset**: Temple University Hospital EEG Corpus (TUEG), the largest public EEG corpus
68
- - **Pretraining Task**: Self-supervised masked EEG patch reconstruction from both time-domain
69
- and frequency-domain EEG signals
70
- - **Model Parameters**: ~4.0M parameters (very compact compared to other foundation models)
71
- - **Fast Convergence**: Achieves decent results in first epoch on downstream tasks,
72
- full convergence within ~10 epochs (vs. ~30 for supervised models like EEGConformer)
73
-
74
- .. rubric:: Macro Components
75
-
76
- - **Patch Encoding Network**: Converts raw EEG patches into embeddings
77
- - **Asymmetric Conditional Positional Encoding (ACPE)**: Generates spatial-temporal positional
78
- embeddings adaptively from input EEG format
79
- - **Criss-Cross Transformer Blocks** (12 layers): Alternates spatial and temporal attention
80
- to learn EEG representations
81
- - **Reconstruction Head**: Reconstructs masked EEG patches during pretraining
82
- - **Task head** (``final_layer``): flatten summary tokens across patches and map to
83
- ``n_outputs``; if ``return_encoder_output=True``, return the encoder features instead.
84
-
85
- The model is highly efficient, requiring only ~318.9M FLOPs on a typical 16-channel, 10-second
86
- EEG recording (significantly lower than full attention baselines).
87
-
88
- .. rubric:: Known Limitations
89
-
90
- - **Data Quality**: TUEG corpus contains "dirty data"; pretraining used crude filtering,
91
- reducing available pre-training data
92
- - **Channel Dependency**: Performance degrades with very sparse electrode setups (e.g., <4 channels)
93
- - **Computational Resources**: While efficient, foundation models have higher deployment
94
- requirements than lightweight models
95
- - **Limited Scaling Exploration**: Future work should explore scaling laws at billion-parameter levels
96
- and integration with large pre-trained vision/language models
97
-
98
- .. important::
99
- **Pre-trained Weights Available**
100
-
101
- This model has pre-trained weights available on the Hugging Face Hub.
102
- You can load them using:
103
-
104
- .. code-block:: python
105
-
106
- from braindecode.models import CBraMod
107
-
108
- # Load pre-trained model from Hugging Face Hub
109
- model = CBraMod.from_pretrained(
110
- "braindecode/cbramod-pretrained", return_encoder_output=True
111
- )
112
-
113
- To push your own trained model to the Hub:
114
-
115
- .. code-block:: python
116
-
117
- # After training your model
118
- model.push_to_hub(
119
- repo_id="username/my-cbramod-model", commit_message="Upload trained CBraMod model"
120
- )
121
-
122
- Requires installing ``braindecode[hug]`` for Hub integration.
123
-
124
- Parameters
125
- ----------
126
- patch_size : int, default=200
127
- Temporal patch size in samples (200 samples = 1 second at 200 Hz).
128
- dim_feedforward : int, default=800
129
- Dimension of the feedforward network in Transformer layers.
130
- n_layer : int, default=12
131
- Number of Transformer layers.
132
- nhead : int, default=8
133
- Number of attention heads.
134
- activation : type[nn.Module], default=nn.GELU
135
- Activation function used in Transformer feedforward layers.
136
- emb_dim : int, default=200
137
- Output embedding dimension.
138
- drop_prob : float, default=0.1
139
- Dropout probability.
140
- return_encoder_output : bool, default=False
141
- If false (default), the features are flattened and passed through a final linear layer
142
- to produce class logits of size ``n_outputs``.
143
- If True, the model returns the encoder output features.
144
-
145
- References
146
- ----------
147
- .. [cbramod] Wang, J., Zhao, S., Luo, Z., Zhou, Y., Jiang, H., Li, S., Li, T., & Pan, G. (2025).
148
- CBraMod: A Criss-Cross Brain Foundation Model for EEG Decoding.
149
- In The Thirteenth International Conference on Learning Representations (ICLR 2025).
150
- https://arxiv.org/abs/2412.07236
151
- """
152
-
153
- def __init__(
154
- self,
155
- n_outputs=None,
156
- n_chans=None,
157
- chs_info=None,
158
- n_times=None,
159
- input_window_seconds=None,
160
- sfreq=None,
161
- patch_size: int = 200,
162
- dim_feedforward: int = 800,
163
- n_layer: int = 12,
164
- nhead: int = 8,
165
- activation: type[nn.Module] = nn.GELU,
166
- emb_dim: int = 200,
167
- channels_kernel_stride_padding_norm: Sequence[
168
- tuple[int, int, int, int, tuple[int, int]]
169
- ] = (
170
- (25, 49, 25, 24, (5, 25)),
171
- (25, 3, 1, 1, (5, 25)),
172
- (25, 3, 1, 1, (5, 25)),
173
- ),
174
- drop_prob: float = 0.1,
175
- return_encoder_output: bool = False,
176
- ):
177
- super().__init__(
178
- n_outputs=n_outputs,
179
- n_chans=n_chans,
180
- chs_info=chs_info,
181
- n_times=n_times,
182
- input_window_seconds=input_window_seconds,
183
- sfreq=sfreq,
184
- )
185
- del n_chans, chs_info, n_times, input_window_seconds, sfreq, n_outputs
186
- self.rearrange = Rearrange(
187
- "batch n_chans (n_patch patch_size) -> batch n_chans n_patch patch_size",
188
- patch_size=patch_size,
189
- )
190
- self.patch_embedding = _PatchEmbedding(
191
- patch_size,
192
- channels_kernel_stride_padding_norm,
193
- drop_prob=drop_prob,
194
- )
195
- d_model = self.patch_embedding.d_model
196
- encoder_layer = CrissCrossTransformerEncoderLayer(
197
- d_model=d_model,
198
- nhead=nhead,
199
- dim_feedforward=dim_feedforward,
200
- batch_first=True,
201
- norm_first=True,
202
- activation=activation,
203
- dropout=drop_prob,
204
- )
205
- self.encoder = _TransformerEncoder(encoder_layer, num_layers=n_layer)
206
- self.proj_out = nn.Sequential(nn.Linear(d_model, emb_dim))
207
-
208
- self._weights_init()
209
-
210
- self.final_layer = (
211
- nn.Identity()
212
- if return_encoder_output
213
- else nn.Sequential(nn.Flatten(), nn.LazyLinear(self.n_outputs))
214
- )
215
-
216
- def _weights_init(self):
217
- for m in self.modules():
218
- if isinstance(m, nn.Linear):
219
- nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
220
- elif isinstance(m, nn.Conv1d):
221
- nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
222
- elif isinstance(m, nn.BatchNorm1d):
223
- nn.init.constant_(m.weight, 1)
224
- nn.init.constant_(m.bias, 0)
225
-
226
- def forward(self, x, mask=None):
227
- x = self.rearrange(x)
228
- patch_emb = self.patch_embedding(x, mask)
229
- feats = self.encoder(patch_emb)
230
- out = self.proj_out(feats)
231
- return self.final_layer(out)
232
-
233
-
234
- class _PatchEmbedding(nn.Module):
235
- def __init__(
236
- self,
237
- patch_size,
238
- channels_kernel_stride_padding_norm,
239
- drop_prob=0.1,
240
- ):
241
- super().__init__()
242
- self.patch_size = patch_size
243
- self.channels_kernel_stride_padding_norm = channels_kernel_stride_padding_norm
244
- self.positional_encoding = nn.Sequential(
245
- nn.Conv2d(
246
- in_channels=self.d_model,
247
- out_channels=self.d_model,
248
- kernel_size=(19, 7),
249
- stride=(1, 1),
250
- padding=(9, 3),
251
- groups=self.d_model,
252
- ),
253
- )
254
- self.mask_encoding = nn.Parameter(torch.zeros(patch_size), requires_grad=False)
255
-
256
- last_channels = 1
257
- proj_in_layers = []
258
- for (
259
- channels,
260
- kernel,
261
- stride,
262
- padding,
263
- norm,
264
- ) in channels_kernel_stride_padding_norm:
265
- proj_in_layers.extend(
266
- [
267
- nn.Conv2d(
268
- in_channels=last_channels,
269
- out_channels=channels,
270
- kernel_size=(1, kernel),
271
- stride=(1, stride),
272
- padding=(0, padding),
273
- ),
274
- nn.GroupNorm(*norm),
275
- nn.GELU(),
276
- ]
277
- )
278
- last_channels = channels
279
- self.proj_in = nn.Sequential(*proj_in_layers)
280
- self.spectral_proj = nn.Sequential(
281
- nn.Linear(patch_size // 2 + 1, self.d_model),
282
- nn.Dropout(drop_prob),
283
- )
284
-
285
- @property
286
- def d_model(self):
287
- last_channels = self.channels_kernel_stride_padding_norm[-1][0]
288
- patch_size = self.patch_size
289
- for _, kernel, stride, padding, _ in self.channels_kernel_stride_padding_norm:
290
- patch_size = int((patch_size + 2 * padding - kernel) / stride + 1)
291
- return last_channels * patch_size
292
-
293
- def forward(self, x, mask=None):
294
- bz, ch_num, patch_num, patch_size = x.shape
295
- if mask is None:
296
- mask_x = x
297
- else:
298
- mask_x = x.clone()
299
- mask_x[mask == 1] = self.mask_encoding
300
-
301
- mask_x = rearrange(mask_x, "b c n p -> b 1 (c n) p")
302
- patch_emb = self.proj_in(mask_x)
303
- patch_emb = rearrange(patch_emb, "b d (c n) p2 -> b c n (d p2)", c=ch_num)
304
-
305
- mask_x = rearrange(mask_x, "b 1 (c n) p -> (b c n) p", c=ch_num)
306
- spectral = torch.fft.rfft(mask_x, dim=-1, norm="forward")
307
- spectral = rearrange(
308
- torch.abs(spectral), "(b c n) p -> b c n p", b=bz, c=ch_num, p=101
309
- )
310
- spectral_emb = self.spectral_proj(spectral)
311
-
312
- patch_emb = patch_emb + spectral_emb
313
-
314
- positional_embedding = self.positional_encoding(
315
- rearrange(patch_emb, "b c n p -> b p c n", p=self.d_model)
316
- ) # d for sanity check
317
- positional_embedding = rearrange(positional_embedding, "b p c n -> b c n p")
318
-
319
- patch_emb = patch_emb + positional_embedding
320
-
321
- return patch_emb
322
-
323
-
324
- class _TransformerEncoder(nn.Module):
325
- def __init__(
326
- self,
327
- encoder_layer,
328
- num_layers,
329
- norm=None,
330
- ):
331
- super().__init__()
332
- self.layers = nn.ModuleList(
333
- [copy.deepcopy(encoder_layer) for i in range(num_layers)]
334
- )
335
- self.num_layers = num_layers
336
- self.norm = norm
337
-
338
- def forward(
339
- self,
340
- src: Tensor,
341
- mask: Optional[Tensor] = None,
342
- src_key_padding_mask: Optional[Tensor] = None,
343
- is_causal: Optional[bool] = None,
344
- ) -> Tensor:
345
- output = src
346
- for mod in self.layers:
347
- output = mod(output, src_mask=mask)
348
- if self.norm is not None:
349
- output = self.norm(output)
350
- return output
@@ -1 +0,0 @@
1
- __version__ = "1.3.1.dev177756420"