Spaces:
Running on Zero
Running on Zero
File size: 3,412 Bytes
093748c | 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 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 | import torch
import torch.nn.functional as F
from setup.configProcessor import select_model_type_from_config, get_in_channels, get_feature_channels, get_dropouts, \
get_out_channels, get_activation, get_dropouts_fcn, get_additional_channels, use_dssp
from misc.enums import ModelType
from machine_learning.Models import GCNConvModel, SAGEConvModel, SAGEConvMLPModel, SAGEConvGATMLPModel
def initialize_model_with_config_params(model_parameters):
model_type = select_model_type_from_config()
classifier = None
match model_type:
case ModelType.GCNCONV:
classifier = GCNConvModel(
in_channels=model_parameters["in_channels"],
feature_channels=model_parameters["features"],
additional_channels=model_parameters["nr_dssp_features"],
out_channels=model_parameters["out_channels"],
dropout=model_parameters["dropout"],
use_additional_channels=model_parameters["use_dssp"]
)
case ModelType.SAGECONV:
classifier = SAGEConvModel(
in_channels=model_parameters["in_channels"],
feature_channels=model_parameters["features"],
additional_channels=model_parameters["nr_dssp_features"],
out_channels=model_parameters["out_channels"],
dropout=model_parameters["dropout"],
activation=eval(model_parameters["activation"]),
use_additional_channels=model_parameters["use_dssp"]
)
case ModelType.SAGECONVMLP:
classifier = SAGEConvMLPModel(
in_channels=model_parameters["in_channels"],
feature_channels=model_parameters["features"],
additional_channels=model_parameters["nr_dssp_features"],
out_channels=model_parameters["out_channels"],
dropout=model_parameters["dropout"],
activation=eval(model_parameters["activation"]),
heads=4, #TODO: Add heads to config?
dropout_fcn=model_parameters["droupout_fcn"],
use_additional_channels=model_parameters["use_dssp"]
)
case ModelType.SAGECONVGATMLP:
classifier = SAGEConvGATMLPModel(
in_channels=model_parameters["in_channels"],
feature_channels=model_parameters["features"],
additional_channels=model_parameters["nr_dssp_features"],
out_channels=model_parameters["out_channels"],
dropout=model_parameters["dropout"],
activation=eval(model_parameters["activation"]),
heads=4, #TODO: Add heads to config?
dropout_fcn=model_parameters["droupout_fcn"],
use_additional_channels=model_parameters["use_dssp"]
)
return classifier
def save_classifier_torch(classifier, model_path):
"""Save pre-trained model"""
torch.save(classifier.state_dict(), model_path)
def load_classifier_torch(model_path, model_parameters):
""" Load pre-saved model """
classifier = initialize_model_with_config_params(model_parameters)
classifier.load_state_dict(torch.load(model_path))
return classifier
def initialize_untrained_model(model_parameters):
classifier = initialize_model_with_config_params(model_parameters)
return classifier
|