Transformers documentation
自定义层和工具
自定义层和工具
此页面列出了库使用的所有自定义层,以及它为模型提供的实用函数。
其中大多数只有在您研究库中模型的代码时才有用。
Pytorch自定义模块
class transformers.Conv1D
< source >( nf nx )
1D-convolutional layer as defined by Radford et al. for OpenAI GPT (and also used in GPT-2).
Basically works like a linear layer but the weights are transposed.
PyTorch帮助函数
transformers.apply_chunking_to_forward
< source >( forward_fn: Callable[..., torch.Tensor] chunk_size: int chunk_dim: int *input_tensors ) → torch.Tensor
Parameters
- forward_fn (
Callable[..., torch.Tensor]
) — The forward function of the model. - chunk_size (
int
) — The chunk size of a chunked tensor:num_chunks = len(input_tensors[0]) / chunk_size
. - chunk_dim (
int
) — The dimension over which theinput_tensors
should be chunked. - input_tensors (
tuple[torch.Tensor]
) — The input tensors offorward_fn
which will be chunked
Returns
torch.Tensor
A tensor with the same shape as the forward_fn
would have given if applied`.
This function chunks the input_tensors
into smaller input tensor parts of size chunk_size
over the dimension
chunk_dim
. It then applies a layer forward_fn
to each chunk independently to save memory.
If the forward_fn
is independent across the chunk_dim
this function will yield the same result as directly
applying forward_fn
to input_tensors
.
Examples:
# rename the usual forward() fn to forward_chunk()
def forward_chunk(self, hidden_states):
hidden_states = self.decoder(hidden_states)
return hidden_states
# implement a chunked forward function
def forward(self, hidden_states):
return apply_chunking_to_forward(self.forward_chunk, self.chunk_size_lm_head, self.seq_len_dim, hidden_states)
transformers.pytorch_utils.prune_linear_layer
< source >( layer: nn.Linear index: torch.LongTensor dim: int = 0 ) → torch.nn.Linear
Prune a linear layer to keep only entries in index.
Used to remove heads.