def get_model(args,
atomref=None,
mean=None,
stddev=None,
train_loader=None,
parallelize=False):
if args.model == 'schnet':
representation = spk.representation.SchNet(args.features,
args.features,
args.interactions,
args.cutoff,
args.num_gaussians,
normalize_filter=True)
atomwise_output = spk.atomistic.Atomwise(
args.features,
aggregation_mode=args.aggregation_mode,
mean=mean,
stddev=stddev,
atomref=atomref,
train_embeddings=True)
model = spk.atomistic.AtomisticModel(representation, atomwise_output)
else:
raise ValueError('Unsupported model class:', args.model)
if parallelize:
model = nn.DataParallel(model)
logging.info("The model you built has: {} parameters".format(
compute_params(model)))
return model
def get_model(args, atomref=None, mean=None, stddev=None, train_loader=None, parallelize=False, mode='train'):
if args.model == 'schnet':
representation = spk.representation.SchNet(args.features, args.features, args.interactions,
args.cutoff, args.num_gaussians)
if args.property == QM9.mu:
atomwise_output = spk.atomistic.DipoleMoment(args.features, predict_magnitude=True, mean=mean,
stddev=stddev)
else:
atomwise_output = spk.atomistic.Atomwise(args.features, aggregation_mode=args.aggregation_mode, mean=mean,
stddev=stddev, atomref=atomref)
model = spk.atomistic.AtomisticModel(representation, atomwise_output)
elif args.model == 'wacsf':
sfmode = ('weighted', 'Behler')[args.behler]
# Convert element strings to atomic charges
elements = frozenset((atomic_numbers[i] for i in sorted(args.elements)))
representation = spk.representation.BehlerSFBlock(args.radial, args.angular, zetas=set(args.zetas),
cutoff_radius=args.cutoff,
centered=args.centered, crossterms=args.crossterms,
mode=sfmode,
elements=elements)
logging.info("Using {:d} {:s}-type SF".format(representation.n_symfuncs, sfmode))
# Standardize representation if requested
if args.standardize and mode == 'train':
if train_loader is None:
raise ValueError("Specification of a training_loader is required to standardize wACSF")
else:
logging.info("Computing symmetry function statistics")
else:
train_loader = None
representation = spk.representation.StandardizeSF(representation, train_loader, cuda=args.cuda)
# Build HDNN model
if args.property == QM9.mu:
atomwise_output = spk.atomistic.ElementalDipoleMoment(representation.n_symfuncs, n_hidden=args.n_nodes,
n_layers=args.n_layers, predict_magnitude=True,
elements=elements)
else:
atomwise_output = spk.atomistic.ElementalAtomwise(representation.n_symfuncs, n_hidden=args.n_nodes,
n_layers=args.n_layers,
aggregation_mode=args.aggregation_mode,
mean=mean, stddev=stddev, atomref=atomref,
elements=elements)
model = spk.atomistic.AtomisticModel(representation, atomwise_output)
else:
raise ValueError('Unknown model class:', args.model)
if parallelize:
model = nn.DataParallel(model)
logging.info("The model you built has: %d parameters" % compute_params(model))
return model
def get_model(representation, output_modules, parallelize=False):
model = AtomisticModel(representation, output_modules)
if parallelize:
model = nn.DataParallel(model)
logging.info("The model you built has: %d parameters" %
compute_params(model))
return model
def get_model(
args, atomref=None, mean=None, stddev=None, train_loader=None, parallelize=False
):
if train_args.cutoff_function == "hard":
cutoff_network = HardCutoff
elif train_args.cutoff_function == "cosine":
cutoff_network = CosineCutoff
elif train_args.cutoff_function == "mollifier":
cutoff_network = MollifierCutoff
if args.model == "schnet":
representation = spk.representation.SchNet(
args.features,
args.features,
args.interactions,
args.cutoff,
args.num_gaussians,
normalize_filter=True,
cutoff_network=cutoff_network,
)
atomwise_output = spk.atomistic.Atomwise(
args.features,
aggregation_mode=args.aggregation_mode,
mean=mean,
stddev=stddev,
atomref=atomref,
train_embeddings=True,
)
model = spk.atomistic.AtomisticModel(representation, atomwise_output)
else:
raise ValueError("Unsupported model class:", args.model)
if parallelize:
model = nn.DataParallel(model)
logging.info("The model you built has: {} parameters".format(compute_params(model)))
return model
def get_model(
args,
atomref=None,
mean=None,
stddev=None,
train_loader=None,
parallelize=False,
mode="train",
):
if args.model == "schnet":
representation = spk.representation.SchNet(
args.features,
args.features,
args.interactions,
args.cutoff,
args.num_gaussians,
)
atomwise_output = spk.atomistic.Energy(
args.features,
mean=mean,
stddev=stddev,
atomref=atomref,
return_force=True,
create_graph=True,
)
model = spk.atomistic.AtomisticModel(representation, atomwise_output)
elif args.model == "wacsf":
sfmode = ("weighted", "Behler")[args.behler]
# Convert element strings to atomic charges
elements = frozenset(
(atomic_numbers[i] for i in sorted(args.elements)))
representation = spk.representation.BehlerSFBlock(
args.radial,
args.angular,
zetas=set(args.zetas),
cutoff_radius=args.cutoff,
centered=args.centered,
crossterms=args.crossterms,
elements=elements,
mode=sfmode,
)
logging.info("Using {:d} {:s}-type SF".format(
representation.n_symfuncs, sfmode))
# Standardize representation if requested
if args.standardize and mode == "train":
if train_loader is None:
raise ValueError(
"Specification of a trainig_loader is required to standardize wACSF"
)
else:
logging.info(
"Computing and standardizing symmetry function statistics")
else:
train_loader = None
representation = spk.representation.StandardizeSF(representation,
train_loader,
cuda=args.cuda)
# Build HDNN model
atomwise_output = spk.atomistic.ElementalEnergy(
representation.n_symfuncs,
n_hidden=args.n_nodes,
n_layers=args.n_layers,
mean=mean,
stddev=stddev,
atomref=atomref,
return_force=True,
create_graph=True,
elements=elements,
)
model = spk.atomistic.AtomisticModel(representation, atomwise_output)
else:
raise ValueError("Unknown model class:", args.model)
if parallelize:
model = nn.DataParallel(model)
logging.info("The model you built has: %d parameters" %
compute_params(model))
return model
def get_model(args,
atomref=None,
mean=None,
stddev=None,
train_loader=None,
parallelize=False,
mode='train'):
if args.model == 'schnet':
if args.pretrained_model:
pretrained_args = json.load(
open(args.pretrained_model + 'args.json', 'r'))
if pretrained_args['interactions'] > args.interactions:
raise ValueError(
'Please set number of interaction layers greater or equal to that of pretrained model'
)
args.features = pretrained_args['features']
args.cutoff = pretrained_args['cutoff']
args.num_gaussians = pretrained_args['num_gaussians']
pretrained = get_pretrained(args.pretrained_model)
if args.edge_updates:
representation = ShrinkSchNet(
args.features,
args.features,
args.interactions,
args.cutoff,
args.num_gaussians,
coupled_interactions=False,
shrink_distances=args.shrink_distances)
else:
representation = ShrinkSchNet(
args.features,
args.features,
args.interactions,
args.cutoff,
args.num_gaussians,
shrink_distances=args.shrink_distances,
edgeupdate_block=None,
coupled_interactions=False)
if args.pretrained_model:
list_of_layers = list(pretrained.interactions.children()
)[:pretrained_args['interactions']]
for child in list_of_layers:
for param in child.parameters():
param.requires_grad = False
# list_of_layers.extend(list(representation.interactions.children())[len(list_of_layers):])
# new_interactions = nn.ModuleList(list_of_layers)
# pretrained.interactions = new_interactions
representation = pretrained
atomwise_output = SCReadout(args.features,
args.features,
args.readout_layers,
mean=mean)
model = spk.atomistic.AtomisticModel(representation, atomwise_output)
print(model)
if parallelize:
model = nn.DataParallel(model)
logging.info("The model you built has: %d parameters" %
compute_params(model))
return model
