File size: 1,202 Bytes
ff0fadf | 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 | from . import Transolver
from . import LSM
from . import FNO
from . import U_Net
from . import Transformer
from . import Factformer
from . import Swin_Transformer
from . import Galerkin_Transformer
from . import GNOT
from . import U_NO
from . import U_FNO
from . import F_FNO
from . import ONO
from . import MWT
from . import GraphSAGE
from . import Graph_UNet
from . import PointNet
from . import DeepONet
from . import MeshGraphNet
from . import RegDGCNN
from . import GFNO
def get_model(args, device):
model_dict = {
"PointNet": PointNet,
"Graph_UNet": Graph_UNet,
"GraphSAGE": GraphSAGE,
"MWT": MWT,
"ONO": ONO,
"F_FNO": F_FNO,
"U_FNO": U_FNO,
"U_NO": U_NO,
"GNOT": GNOT,
"Galerkin_Transformer": Galerkin_Transformer,
"Swin_Transformer": Swin_Transformer,
"Factformer": Factformer,
"Transformer": Transformer,
"U_Net": U_Net,
"FNO": FNO,
"Transolver": Transolver,
"LSM": LSM,
"DeepONet": DeepONet,
"MeshGraphNet": MeshGraphNet,
"RegDGCNN": RegDGCNN,
"GFNO": GFNO,
}
return model_dict[args.model].Model(args, device)
|