File size: 1,540 Bytes
6177cee | 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 | """
(c) Adaptation of the code from https://github.com/THUDM/GraphMAE
"""
from .edcoder import PreModel
def build_model(args):
num_heads = args.num_heads
num_out_heads = args.num_out_heads
num_hidden = args.num_hidden
num_layers = args.num_layers
residual = args.residual
attn_drop = args.attn_drop
in_drop = args.in_drop
norm = args.norm
negative_slope = args.negative_slope
encoder_type = args.encoder
decoder_type = args.decoder
mask_rate = args.mask_rate
drop_edge_rate = args.drop_edge_rate
replace_rate = args.replace_rate
batchnorm = args.batchnorm
activation = args.activation
loss_fn = args.loss_fn
alpha_l = args.alpha_l
concat_hidden = args.concat_hidden
num_features = args.num_features
num_edge_features = args.num_edge_features
model = PreModel(
in_dim=int(num_features),
edge_in_dim=int(num_edge_features),
num_hidden=int(num_hidden),
num_layers=num_layers,
nhead=num_heads,
nhead_out=num_out_heads,
activation=activation,
feat_drop=in_drop,
attn_drop=attn_drop,
negative_slope=negative_slope,
residual=residual,
encoder_type=encoder_type,
decoder_type=decoder_type,
mask_rate=mask_rate,
norm=norm,
loss_fn=loss_fn,
drop_edge_rate=drop_edge_rate,
replace_rate=replace_rate,
alpha_l=alpha_l,
concat_hidden=concat_hidden,
batchnorm=batchnorm,
)
return model
|