File size: 1,421 Bytes
35f60b9 a532d9f 35f60b9 a532d9f 35f60b9 a532d9f 35f60b9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 | 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 |