Diffusers documentation
RFD4ProteinaAutoencoder
RFD4ProteinaAutoencoder
RFD4-Proteína 的自编码器。encode 将 atom37 坐标映射为每个 token 的局部潜在变量，而 decode 则将潜在变量和骨架位置还原为 atom37 结构和残基身份。它在流水线中注册为 vae。
from diffusers import RFD4ProteinaAutoencoder
vae = RFD4ProteinaAutoencoder.from_pretrained("nvidia/RFD4-Proteina-Diffusers", subfolder="vae")RFD4ProteinaAutoencoder
class diffusers.RFD4ProteinaAutoencoder
< source >( atom37_mask: collections.abc.Sequence[collections.abc.Sequence[bool]]allowed_residue_types: collections.abc.Sequence[collections.abc.Sequence[bool]]latent_dim: int = 32n_residue_types: int = 147num_token_types: int = 7nlayers: int = 12token_dim: int = 768nheads: int = 12dim_cond: int = 128pair_repr_dim: int = 256token_type_emb_dim: int = 32encoder_feats_seq: tuple[str, ...] = ('token_type', 'x1_a37coors_nm', 'x1_a37coors_nm_rel', 'chain_idx_seq', 'incident_bond_type_counts', 'atom37_element_charge_embedding')decoder_feats_seq: tuple[str, ...] = ('ca_coors_nm', 'z_latent_seq', 'token_type', 'cond_residue_type') )
参数
- latent_dim (
int，默认为32) — 每个 token 的局部隐向量的宽度。 - n_residue_types (
int，默认为147) — 残基词汇表大小。 - num_token_types (
int，默认为7) — 标记类型词汇表的大小：蛋白质、RNA、DNA、原子化、未知、聚糖和非标准氨基酸。 - allowed_residue_types (
list[list[bool]]) — 一个(num_token_types, n_residue_types)的残基类别表，表示每种 token 类型可能解码到的残基类别。 - nlayers (
int，默认为12) — 编码器和解码器中各自的 Trunk 层数。 - token_dim (
int，默认为768) — Token 表示的宽度。 - nheads（
int，默认为12）— 每层的注意力头数。 - dim_cond（
int，默认为128）—（未使用的、全零的）条件向量的宽度。 - pair_repr_dim (
int，默认为256) — 用于构建主干层的配对宽度。发布的自编码器不使用配对偏置，因此该参数不包含任何可学习参数。 - token_type_emb_dim (
int，默认为32) — Token 类型嵌入的宽度。 - encoder_feats_seq, decoder_feats_seq (
list[str]) — 投影到编码器和解码器输入中的特征。 - atom37_mask (
list[list[bool]]) — 每种残基类型对应的有效 atom37 槽位，形状为(n_residue_types, 37)。
配对的分子编码器和 atom37 解码器。
提供给该模型以及由该模型返回的坐标单位为纳米。
encode
< source >( coords_nm: Tensorpadding_mask: Tensortoken_type: typing.Optional[torch.Tensor] = Noneatom37_mask: typing.Optional[torch.Tensor] = Noneconditioning: dict[str, typing.Any] | None = Nonegenerator: typing.Union[torch.Generator, list[torch.Generator], NoneType] = Nonereturn_dict: bool = True ) → RFD4ProteinaEncoderOutput or tuple
参数
- coords_nm (
torch.Tensor) —(batch, tokens, 37, 3)以纳米为单位的 atom37 坐标。 - padding_mask (
torch.Tensor) —(batch, tokens)，其中真实 token 的位置为 true，填充位置为 false。 - token_type (
torch.Tensor，可选) —(batch, tokens)的 token 类型编码。 - atom37_mask (
torch.Tensor, 可选) —(batch, tokens, 37)，在原子槽位被使用时为 true。 - conditioning (
dict[str, Any]，可选) — 其余的拓扑和条件张量，位于预处理器输出的键下。 - generator (
torch.Generator或list[torch.Generator]，可选) — 每个批次行对应一个生成器，用于生成独立的随机潜在样本。若未提供，则采样将使用全局随机流。 - return_dict (
bool，默认为True) — 返回带类型的输出，或返回其字段的元组。
返回值
RFD4ProteinaEncoderOutput 或 tuple
形状为 (batch, tokens, latent_dim) 的潜在、均值和对数尺度张量。
将准备好的坐标和条件编码为采样的逐 token 局部潜在变量。
decode
< source >( z_latent: Tensorca_coors_nm: Tensorpadding_mask: Tensortoken_type: typing.Optional[torch.Tensor] = Noneatom37_mask: typing.Optional[torch.Tensor] = Noneconditioning: dict[str, typing.Any] | None = Nonereturn_dict: bool = True ) → RFD4ProteinaDecoderOutput or tuple
参数
- z_latent (
torch.Tensor) — 每个 token 的局部潜在表示，形状为(batch, tokens, latent_dim)。 - ca_coors_nm (
torch.Tensor) —(batch, tokens, 3)以纳米为单位的 backbone CA 位置。 - padding_mask (
torch.Tensor) —(batch, tokens)，其中真实 token 的位置为 true，填充位置为 false。 - token_type (
torch.Tensor，可选) —(batch, tokens)的 token 类型编码。 - atom37_mask (
torch.Tensor, 可选) —(batch, tokens, 37)，在原子槽位被使用时为 true。 - conditioning (
dict[str, Any]，可选) — 其余的拓扑和条件张量， 位于预处理器输出的键下。序列条件将从此处读取。 - return_dict (
bool，默认为True) — 返回带类型的输出，或返回其字段的元组。
返回值
RFD4ProteinaDecoderOutput 或 tuple
以纳米为单位的坐标 (batch, tokens, 37, 3)、残基类型 (batch, tokens)、发射掩码 (batch, tokens) 和原子掩码 (batch, tokens, 37)。未发射原子的坐标为零。
将每个 token 的潜在变量和主干位置解码为 atom37 结构。
forward
< source >( z_latent: Tensorca_coors_nm: Tensorpadding_mask: Tensortoken_type: typing.Optional[torch.Tensor] = Noneatom37_mask: typing.Optional[torch.Tensor] = Noneconditioning: dict[str, typing.Any] | None = Nonereturn_dict: bool = True ) → RFD4ProteinaDecoderOutput or tuple
参数
- z_latent (
torch.Tensor) — 每个 token 的局部潜在表示，形状为(batch, tokens, latent_dim)。 - ca_coors_nm (
torch.Tensor) —(batch, tokens, 3)以纳米为单位的 backbone CA 位置。 - padding_mask (
torch.Tensor) —(batch, tokens)，其中真实 token 的位置为 true，填充位置为 false。 - token_type (
torch.Tensor，可选) —(batch, tokens)的 token 类型编码。 - atom37_mask (
torch.Tensor, 可选) —(batch, tokens, 37)，在原子槽位被使用时为 true。 - conditioning (
dict[str, Any]，可选) — 其余的拓扑和条件张量。 - return_dict (
bool，默认为True) — 返回带类型的输出，或返回其字段的元组。
返回值
RFD4ProteinaDecoderOutput 或 tuple
解码后的坐标、残基类型以及原子/发射掩码。
解码每个 token 的潜在表示和主干位置；等效于调用 decode()。
RFD4ProteinaEncoderOutput
class diffusers.models.autoencoders.autoencoder_rfd4_proteina.RFD4ProteinaEncoderOutput
< source >( latent: Tensormean: typing.Optional[torch.Tensor] = Nonelog_scale: typing.Optional[torch.Tensor] = None )
采样的潜在变量和高斯参数，每个的形状均为 (batch, tokens, latent_dim)。
RFD4ProteinaDecoderOutput
class diffusers.models.autoencoders.autoencoder_rfd4_proteina.RFD4ProteinaDecoderOutput
< source >( coordinates: Tensorresidue_type: typing.Optional[torch.Tensor] = Noneemit_mask: typing.Optional[torch.Tensor] = Noneatom_mask: typing.Optional[torch.Tensor] = None )
解码后的分子坐标和离散 atom37 拓扑结构。