Linear#
- class cuequivariance_torch.Linear(
- irreps_in: Irreps,
- irreps_out: Irreps,
- *,
- layout: IrrepsLayout | None = None,
- layout_in: IrrepsLayout | None = None,
- layout_out: IrrepsLayout | None = None,
- shared_weights: bool = True,
- internal_weights: bool = None,
- device: device | None = None,
- dtype: dtype | None = None,
- math_dtype: dtype | None = None,
- use_fallback: bool | None = None,
一个表示等变线性层的类。
- 参数:
irreps_in (Irreps) – 输入不可约表示。
irreps_out (Irreps) – 输出不可约表示。
layout (IrrepsLayout, 可选) – 不可约表示的布局,默认为
cue.mul_ir
。这是 e3nn 库中使用的布局。shared_weights (bool, 可选) – 是否使用共享权重,默认为 True。
internal_weights (bool, 可选) – 是否使用内部权重,如果 shared_weights 为 True,则默认为 True,否则为 False。
use_fallback (bool, 可选) – 如果为 None (默认值),则在可用时使用 CUDA 内核。如果为 False,将使用 CUDA 内核,如果不可用则会引发异常。如果为 True,则无论 CUDA 内核是否可用,都使用 PyTorch 后备方法。
前向传播
- forward( ) Tensor #
线性层的前向传播。
- 参数:
x (torch.Tensor) – 输入张量。
weight (torch.Tensor, 可选) – 权重张量。如果为 None,则使用内部权重张量。
- 返回值:
应用线性变换后的输出张量。
- 返回类型:
- 抛出:
ValueError – 如果使用内部权重且 weight 不为 None,或者如果使用共享权重且 weight 不是 1D 张量,或者如果未使用共享权重且 weight 不是 2D 张量。