|
import torch
|
|
import gmm_transformer as gmm_model
|
|
|
|
def load_model(
|
|
n_components = 6,
|
|
hidden_d = 24 * 4,
|
|
out_d = 24,
|
|
n_heads = 4,
|
|
mlp_ratio = 8,
|
|
n_blocks = 6,
|
|
encoder_path = r'_encoder_25_4537398.pth',
|
|
path_para = r'_embedding_25_4537398.pth',
|
|
path_token = r'_emb_empty_token_25_4537398.pth',
|
|
random_sample_num = None
|
|
):
|
|
|
|
chw = (1, random_sample_num, 25)
|
|
|
|
|
|
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
|
|
|
|
|
encoder = gmm_model.ViT_encodernopara(chw, hidden_d, out_d, n_heads, mlp_ratio, n_blocks).to(device)
|
|
_model_scale = sum(p.numel() for p in encoder.parameters() if p.requires_grad)
|
|
print('Number of parameters of encoder:', _model_scale)
|
|
|
|
|
|
encoder.load_state_dict(torch.load(encoder_path, map_location=device))
|
|
|
|
state_dict_para = torch.load(path_para, map_location=device)
|
|
state_dict_token = torch.load(path_token, map_location=device)
|
|
|
|
return encoder, state_dict_para, state_dict_token
|
|
|