atomgen.models package

Models module for the AtomGen library.

This module contains the model classes and functions for training and inference.

Submodules

• atomgen.models.configuration_atomformer module
  • AtomformerConfig
    • AtomformerConfig.model_type
• atomgen.models.modeling_atomformer module
  • AtomFormerForSystemClassification
    • AtomFormerForSystemClassification.forward()
  • AtomformerEncoder
    • AtomformerEncoder.forward()
  • AtomformerForCoordinateAM
    • AtomformerForCoordinateAM.forward()
  • AtomformerForMaskedAM
    • AtomformerForMaskedAM.forward()
  • AtomformerModel
    • AtomformerModel.forward()
  • AtomformerPreTrainedModel
    • AtomformerPreTrainedModel.base_model_prefix
    • AtomformerPreTrainedModel.config_class
    • AtomformerPreTrainedModel.supports_gradient_checkpointing
  • GaussianLayer
    • GaussianLayer.forward()
  • InitialStructure2RelaxedEnergy
    • InitialStructure2RelaxedEnergy.forward()
  • InitialStructure2RelaxedStructure
    • InitialStructure2RelaxedStructure.forward()
  • InitialStructure2RelaxedStructureAndEnergy
    • InitialStructure2RelaxedStructureAndEnergy.forward()
  • ParallelBlock
    • ParallelBlock.forward()
  • Structure2Energy
    • Structure2Energy.forward()
  • Structure2EnergyAndForces
    • Structure2EnergyAndForces.forward()
  • Structure2Forces
    • Structure2Forces.forward()
  • Structure2TotalEnergyAndForces
    • Structure2TotalEnergyAndForces.forward()
• atomgen.models.schnet module
  • SchNetConfig
    • SchNetConfig.vocab_size
    • SchNetConfig.hidden_channels
    • SchNetConfig.num_filters
    • SchNetConfig.num_interactions
    • SchNetConfig.num_gaussians
    • SchNetConfig.cutoff
    • SchNetConfig.interaction_graph
    • SchNetConfig.max_num_neighbors
    • SchNetConfig.readout
    • SchNetConfig.dipole
    • SchNetConfig.mean
    • SchNetConfig.std
    • SchNetConfig.atomref
    • SchNetConfig.mask_token_id
    • SchNetConfig.pad_token_id
    • SchNetConfig.bos_token_id
    • SchNetConfig.eos_token_id
  • SchNetModel
    • SchNetModel.forward()
  • SchNetPreTrainedModel
    • SchNetPreTrainedModel.base_model_prefix
    • SchNetPreTrainedModel.config_class
    • SchNetPreTrainedModel.supports_gradient_checkpointing
• atomgen.models.tokengt module
  • InitialStructure2RelaxedEnergy
    • InitialStructure2RelaxedEnergy.forward()
  • InitialStructure2RelaxedStructure
    • InitialStructure2RelaxedStructure.forward()
  • InitialStructure2RelaxedStructureAndEnergy
    • InitialStructure2RelaxedStructureAndEnergy.forward()
  • ParallelBlock
    • ParallelBlock.forward()
  • Structure2EnergyAndForces
    • Structure2EnergyAndForces.forward()
  • TransformerConfig
  • TransformerEncoder
    • TransformerEncoder.forward()
  • TransformerForCoordinateAM
    • TransformerForCoordinateAM.forward()
  • TransformerForMaskedAM
    • TransformerForMaskedAM.forward()
  • TransformerModel
    • TransformerModel.forward()
  • TransformerPreTrainedModel
    • TransformerPreTrainedModel.base_model_prefix
    • TransformerPreTrainedModel.config_class
    • TransformerPreTrainedModel.supports_gradient_checkpointing