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