| from functools import partial |
| import torch.nn as nn |
|
|
| from models.mae import MultiMaskedAutoencoderConv1d |
| from models.classification import ClassificationMultiConv1d, ClassificationMultiConv1dMamba |
| from models.transformer_modules import RMSNorm |
|
|
| def convmae1d_multi(**kwargs): |
| model = MultiMaskedAutoencoderConv1d( |
| embed_dim=[32, 64, 128, 384], depth=[4, 4], specific_depth=[2, 2], num_heads=8, |
| decoder_embed_dim=64, decoder_depth=2, decoder_num_heads=4, |
| fusion_method='cross_attn', fusion_heads=8, fusion_depth=4, |
| mlp_ratio=4, norm_layer=partial(RMSNorm, eps=1e-8), pos_enc='rope', **kwargs) |
| return model |
|
|
| def classification_multi(**kwargs): |
| model = ClassificationMultiConv1d( |
| embed_dim=[32, 64, 128, 384], depth=[4, 4], specific_depth=[2, 2], num_heads=8, |
| fusion_method='cross_attn', fusion_heads=8, fusion_depth=4, |
| mlp_ratio=4, norm_layer=partial(RMSNorm, eps=1e-8), pos_enc='rope', **kwargs) |
| return model |
|
|
| def classification_multi_full_night(**kwargs): |
| model = ClassificationMultiConv1dMamba( |
| embed_dim=[32, 64, 128, 384], depth=[4, 4], specific_depth=[2, 2], num_heads=8, |
| fusion_method='cross_attn', fusion_heads=8, fusion_depth=4, |
| seq_dim=512, seq_depth=1, |
| mlp_ratio=4, norm_layer=partial(RMSNorm, eps=1e-8), pos_enc='rope', **kwargs) |
| return model |
