scgpt.model package

Submodules

scgpt.model.dsbn module

class scgpt.model.dsbn.DomainSpecificBatchNorm1d(num_features: int, num_domains: int, eps: float = 1e-05, momentum: float = 0.1, affine: bool = True, track_running_stats: bool = True)

Bases: _DomainSpecificBatchNorm

property bn_handle: Module

training: bool

class scgpt.model.dsbn.DomainSpecificBatchNorm2d(num_features: int, num_domains: int, eps: float = 1e-05, momentum: float = 0.1, affine: bool = True, track_running_stats: bool = True)

Bases: _DomainSpecificBatchNorm

property bn_handle: Module

training: bool

scgpt.model.generation_model module

class scgpt.model.generation_model.ClsDecoder(d_model: int, n_cls: int, nlayers: int = 3, activation: callable = <class 'torch.nn.modules.activation.ReLU'>)

Bases: Module

Decoder for classification task.

forward(x: Tensor) → Tensor

Parameters:

x – Tensor, shape [batch_size, embsize]

training: bool

class scgpt.model.generation_model.GeneEncoder(num_embeddings: int, embedding_dim: int, padding_idx: Optional[int] = None)

Bases: Module

forward(x: Tensor) → Tensor

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

training: bool

class scgpt.model.generation_model.PositionalEncoding(d_model: int, dropout: float = 0.1, max_len: int = 5000)

Bases: Module

forward(x: Tensor) → Tensor

Parameters:

x – Tensor, shape [seq_len, batch_size, embedding_dim]

training: bool

class scgpt.model.generation_model.Similarity(temp)

Bases: Module

Dot product or cosine similarity

forward(x, y)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

training: bool

class scgpt.model.generation_model.TransformerGenerator(ntoken: int, d_model: int, nhead: int, d_hid: int, nlayers: int, nlayers_cls: int, n_cls: int, vocab: Any, dropout: float = 0.5, pad_token: str = '<pad>', pad_value: int = 0, pert_pad_id: int = 2, do_mvc: bool = False, domain_spec_batchnorm: Union[bool, str] = False, cell_emb_style: str = 'cls', mvc_decoder_style: str = 'inner product', ecs_threshold: float = 0.3, explicit_zero_prob: bool = False, use_fast_transformer: bool = False, fast_transformer_backend: str = 'flash', pre_norm: bool = False)

Bases: Module

encode_batch(src: Tensor, values: Tensor, src_key_padding_mask: Tensor, batch_size: int, output_to_cpu: bool = True) → Tensor

Parameters:

• src – Tensor, shape [N, seq_len]

• values – Tensor, shape [N, seq_len]

• src_key_padding_mask – Tensor, shape [N, seq_len]

Returns:

output Tensor of shape [N, seq_len, embsize]

forward(src: Tensor, values: Tensor, input_pert_flags: Tensor, src_key_padding_mask: Tensor, CLS: bool = False, CCE: bool = False, MVC: bool = False, ECS: bool = False, do_sample: bool = False) → Mapping[str, Tensor]

Parameters:

• src (Tensor) – token ids, shape [batch_size, seq_len]

• values (Tensor) – token values, shape [batch_size, seq_len]

• src_key_padding_mask (Tensor) – mask for src, shape [batch_size, seq_len]

• CLS (bool) – if True, return the celltype classification objective (CLS) output

• CCE (bool) – if True, return the contrastive cell embedding objective (CCE) output

• MVC (bool) – if True, return the masked value prediction for cell embedding MVC output

• ECS (bool) – if True, return the elastic cell similarity objective (ECS) output.

Returns:

dict of output Tensors.

init_weights() → None

pred_perturb(batch_data, include_zero_gene='batch-wise', gene_ids=None, amp=True) → Tensor

Parameters:

batch_data – a dictionary of input data with keys.

Returns:

output Tensor of shape [N, seq_len]

training: bool

scgpt.model.generation_model.generate_square_subsequent_mask(sz: int) → Tensor

Generates an upper-triangular matrix of -inf, with zeros on diag.

scgpt.model.grad_reverse module

class scgpt.model.grad_reverse.GradReverse(*args, **kwargs)

Bases: Function

static backward(ctx, grad_output: Tensor) → Tensor

Defines a formula for differentiating the operation with backward mode automatic differentiation (alias to the vjp function).

This function is to be overridden by all subclasses.

It must accept a context ctx as the first argument, followed by as many outputs as the forward() returned (None will be passed in for non tensor outputs of the forward function), and it should return as many tensors, as there were inputs to forward(). Each argument is the gradient w.r.t the given output, and each returned value should be the gradient w.r.t. the corresponding input. If an input is not a Tensor or is a Tensor not requiring grads, you can just pass None as a gradient for that input.

The context can be used to retrieve tensors saved during the forward pass. It also has an attribute ctx.needs_input_grad as a tuple of booleans representing whether each input needs gradient. E.g., backward() will have ctx.needs_input_grad[0] = True if the first input to forward() needs gradient computated w.r.t. the output.

static forward(ctx, x: Tensor, lambd: float) → Tensor

Performs the operation.

This function is to be overridden by all subclasses.

It must accept a context ctx as the first argument, followed by any number of arguments (tensors or other types).

The context can be used to store arbitrary data that can be then retrieved during the backward pass. Tensors should not be stored directly on ctx (though this is not currently enforced for backward compatibility). Instead, tensors should be saved either with ctx.save_for_backward() if they are intended to be used in backward (equivalently, vjp) or ctx.save_for_forward() if they are intended to be used for in jvp.

scgpt.model.grad_reverse.grad_reverse(x: Tensor, lambd: float = 1.0) → Tensor

scgpt.model.model module

class scgpt.model.model.AdversarialDiscriminator(d_model: int, n_cls: int, nlayers: int = 3, activation: callable = <class 'torch.nn.modules.activation.LeakyReLU'>, reverse_grad: bool = False)

Bases: Module

Discriminator for the adversarial training for batch correction.

forward(x: Tensor) → Tensor

Parameters:

x – Tensor, shape [batch_size, embsize]

training: bool

class scgpt.model.model.BatchLabelEncoder(num_embeddings: int, embedding_dim: int, padding_idx: Optional[int] = None)

Bases: Module

forward(x: Tensor) → Tensor

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

training: bool

class scgpt.model.model.CategoryValueEncoder(num_embeddings: int, embedding_dim: int, padding_idx: Optional[int] = None)

Bases: Module

forward(x: Tensor) → Tensor

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

training: bool

class scgpt.model.model.ClsDecoder(d_model: int, n_cls: int, nlayers: int = 3, activation: callable = <class 'torch.nn.modules.activation.ReLU'>)

Bases: Module

Decoder for classification task.

forward(x: Tensor) → Tensor

Parameters:

x – Tensor, shape [batch_size, embsize]

training: bool

class scgpt.model.model.ContinuousValueEncoder(d_model: int, dropout: float = 0.1, max_value: int = 512)

Bases: Module

Encode real number values to a vector using neural nets projection.

forward(x: Tensor) → Tensor

Parameters:

x – Tensor, shape [batch_size, seq_len]

training: bool

class scgpt.model.model.ExprDecoder(d_model: int, explicit_zero_prob: bool = False, use_batch_labels: bool = False)

Bases: Module

forward(x: Tensor) → Dict[str, Tensor]

x is the output of the transformer, (batch, seq_len, d_model)

training: bool

class scgpt.model.model.FastTransformerEncoderWrapper(d_model: int, nhead: int, d_hid: int, nlayers: int, dropout: float = 0.5)

Bases: Module

static build_fast_transformer_encoder(d_model: int, nhead: int, d_hid: int, nlayers: int, dropout: float) → Module

static build_length_mask(src: Tensor, src_key_padding_mask: BoolTensor) → LengthMask

forward(src: Tensor, src_key_padding_mask: BoolTensor) → Tensor

Parameters:

• src – Tensor, shape [N, seq_len, embsize]

• src_key_padding_mask – Tensor, shape [N, seq_len]

Returns:

output Tensor of shape [N, seq_len, embsize]

training: bool

class scgpt.model.model.FlashTransformerEncoderLayer(d_model, nhead, dim_feedforward=2048, dropout=0.1, activation='relu', layer_norm_eps=1e-05, batch_first=True, device=None, dtype=None, norm_scheme='post')

Bases: Module

TransformerEncoderLayer is made up of self-attn and feedforward network. The class is modified from torch.nn.TransformerEncoderLayer to support the FlashAttention.

Parameters:

• d_model – the number of expected features in the input (required).

• nhead – the number of heads in the multiheadattention models (required).

• dim_feedforward – the dimension of the feedforward network model (default=2048).

• dropout – the dropout value (default=0.1).

• activation – the activation function of intermediate layer, relu or gelu (default=relu).

• layer_norm_eps – the eps value in layer normalization components (default=1e-5).

• batch_first – If True, then the input and output tensors are provided as (batch, seq, feature). Default: False.

Examples::

>>> encoder_layer = nn.TransformerEncoderLayer(d_model=512, nhead=8)
>>> src = torch.rand(10, 32, 512)
>>> out = encoder_layer(src)

Alternatively, when batch_first is True:

>>> encoder_layer = nn.TransformerEncoderLayer(d_model=512, nhead=8, batch_first=True)
>>> src = torch.rand(32, 10, 512)
>>> out = encoder_layer(src)

forward(src: Tensor, src_mask: Optional[Tensor] = None, src_key_padding_mask: Optional[Tensor] = None, **kwargs) → Tensor

Pass the input through the encoder layer.

Parameters:

• src – the sequence to the encoder layer (required).

• src_mask – the mask for the src sequence (optional).

• src_key_padding_mask – the mask for the src keys per batch (optional).

Shape:

see the docs in Transformer class.

training: bool

class scgpt.model.model.GeneEncoder(num_embeddings: int, embedding_dim: int, padding_idx: Optional[int] = None)

Bases: Module

forward(x: Tensor) → Tensor

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

training: bool

class scgpt.model.model.MVCDecoder(d_model: int, arch_style: str = 'inner product', query_activation: ~torch.nn.modules.module.Module = <class 'torch.nn.modules.activation.Sigmoid'>, hidden_activation: ~torch.nn.modules.module.Module = <class 'torch.nn.modules.activation.PReLU'>, explicit_zero_prob: bool = False, use_batch_labels: bool = False)

Bases: Module

Decoder for the masked value prediction for cell embeddings.

forward(cell_emb: Tensor, gene_embs: Tensor) → Union[Tensor, Dict[str, Tensor]]

Parameters:

• cell_emb – Tensor, shape (batch, embsize=d_model)

• gene_embs – Tensor, shape (batch, seq_len, embsize=d_model)

training: bool

class scgpt.model.model.PositionalEncoding(d_model: int, dropout: float = 0.1, max_len: int = 5000)

Bases: Module

forward(x: Tensor) → Tensor

Parameters:

x – Tensor, shape [seq_len, batch_size, embedding_dim]

training: bool

class scgpt.model.model.Similarity(temp)

Bases: Module

Dot product or cosine similarity

forward(x, y)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

training: bool

class scgpt.model.model.TransformerModel(ntoken: int, d_model: int, nhead: int, d_hid: int, nlayers: int, nlayers_cls: int = 3, n_cls: int = 1, vocab: Optional[Any] = None, dropout: float = 0.5, pad_token: str = '<pad>', pad_value: int = 0, do_mvc: bool = False, do_dab: bool = False, use_batch_labels: bool = False, num_batch_labels: Optional[int] = None, domain_spec_batchnorm: Union[bool, str] = False, input_emb_style: str = 'continuous', n_input_bins: Optional[int] = None, cell_emb_style: str = 'cls', mvc_decoder_style: str = 'inner product', ecs_threshold: float = 0.3, explicit_zero_prob: bool = False, use_fast_transformer: bool = False, fast_transformer_backend: str = 'flash', pre_norm: bool = False)

Bases: Module

encode_batch(src: Tensor, values: Tensor, src_key_padding_mask: Tensor, batch_size: int, batch_labels: Optional[Tensor] = None, output_to_cpu: bool = True, time_step: Optional[int] = None, return_np: bool = False) → Tensor

Parameters:

• src (Tensor) – shape [N, seq_len]

• values (Tensor) – shape [N, seq_len]

• src_key_padding_mask (Tensor) – shape [N, seq_len]

• batch_size (int) – batch size for encoding

• batch_labels (Tensor) – shape [N, n_batch_labels]

• output_to_cpu (bool) – whether to move the output to cpu

• time_step (int) – the time step index in the transformer output to return. The time step is along the second dimenstion. If None, return all.

• return_np (bool) – whether to return numpy array

Returns:

output Tensor of shape [N, seq_len, embsize]

forward(src: Tensor, values: Tensor, src_key_padding_mask: Tensor, batch_labels: Optional[Tensor] = None, CLS: bool = False, CCE: bool = False, MVC: bool = False, ECS: bool = False, do_sample: bool = False) → Mapping[str, Tensor]

Parameters:

• src (Tensor) – token ids, shape [batch_size, seq_len]

• values (Tensor) – token values, shape [batch_size, seq_len]

• src_key_padding_mask (Tensor) – mask for src, shape [batch_size, seq_len]

• batch_labels (Tensor) – batch labels, shape [batch_size]

• CLS (bool) – if True, return the celltype classification objective (CLS) output

• CCE (bool) – if True, return the contrastive cell embedding objective (CCE) output

• MVC (bool) – if True, return the masked value prediction for cell embedding MVC output

• ECS (bool) – if True, return the elastic cell similarity objective (ECS) output.

Returns:

dict of output Tensors.

generate(cell_emb: Tensor, src: Tensor, values: Optional[Tensor] = None, src_key_padding_mask: Optional[Tensor] = None, gen_iters: int = 1, batch_labels: Optional[Tensor] = None) → Tensor

Parameters:

• cell_emb (Tensor) – shape (batch, embsize)

• src (Tensor) – shape (batch, seq_len)

• values (Tensor) – shape (batch, seq_len), optional

• src_key_padding_mask (Tensor) – shape (batch, seq_len), optional

• gen_iters (int) – number of generation iterations

• batch_labels (Tensor) – shape (batch,), optional

init_weights() → None

training: bool

scgpt.model.model.generate_square_subsequent_mask(sz: int) → Tensor

Generates an upper-triangular matrix of -inf, with zeros on diag.

Module contents