torch_geometric.nn.models.ViSNet

class ViSNet(lmax: int = 1, vecnorm_type: Optional[str] = None, trainable_vecnorm: bool = False, num_heads: int = 8, num_layers: int = 6, hidden_channels: int = 128, num_rbf: int = 32, trainable_rbf: bool = False, max_z: int = 100, cutoff: float = 5.0, max_num_neighbors: int = 32, vertex: bool = False, atomref: Optional[Tensor] = None, reduce_op: str = 'sum', mean: float = 0.0, std: float = 1.0, derivative: bool = False)[source]

Bases: Module

一个模块，实现了来自“通过等变向量-标量交互消息传递增强分子的几何表示”论文的等变向量-标量交互图神经网络（ViSNet）。

Parameters:

  • lmax (int, optional) – 球谐函数的最大阶数。 (默认值: 1)

  • vecnorm_type (str, optional) – 应用于向量的归一化类型。(默认值: None)

  • trainable_vecnorm (bool, optional) – 归一化权重是否可训练。（默认值：False）

  • num_heads (int, optional) – 注意力头的数量。 (默认: 8)

  • num_layers (int, optional) – 网络中的层数。 (默认值: 6)

  • hidden_channels (int, optional) – 节点嵌入中的隐藏通道数。（默认值：128）

  • num_rbf (int, optional) – 径向基函数的数量。 (默认: 32)

  • trainable_rbf (bool, optional) – 径向基函数参数是否可训练。(默认: False)

  • max_z (int, optional) – 最大原子序数。 (默认值: 100)

  • cutoff (float, optional) – 截止距离。（默认值：5.0）

  • max_num_neighbors (int, optional) – 每个原子考虑的最大邻居数。(默认值: 32)

  • 顶点 (bool, 可选) – 是否使用顶点几何特征。 (默认: False)

  • atomref (torch.Tensor, optional) – 一个原子参考值的张量，如果未提供则为 None。（默认值：None）

  • reduce_op (str, optional) – 要应用的归约操作类型 ("sum", "mean")。(默认值: "sum")

  • mean (float, optional) – 输出分布的均值。 (默认值: 0.0)

  • std (float, optional) – 输出分布的标准差。（默认值：1.0）

  • derivative (bool, optional) – 是否计算输出相对于位置的导数。(默认: False)

forward(z: Tensor, pos: Tensor, batch: Tensor) → Tuple[Tensor, Optional[Tensor]][source]

计算一批分子的能量或属性（力）。

Parameters:
Returns:

每个分子的能量或属性。 dy (torch.Tensor, 可选): 能量的负导数。

Return type:

y (torch.Tensor)

reset_parameters()[source]

重置模块的参数。