Skip to main content
Join the official 2019 Python Developers SurveyStart the survey!

A pytorch implementation of transformer encoder

Project description

Transformer Encoder

This package provides an easy-to-use interface of transformer encoder.


Requirements: python(>=3.5), pytorch(>=1.0.0)

Install from pypi:

pip install tfencoder

Or from Github for the latest version:

pip install git+


TFEncoder(n_layers, d_model, d_ff, n_heads, dropout)

  • n_layers: number of stacked layers of encoder
  • d_model: dimension of each word vector
  • d_ff: hidden dimension of feed forward layer
  • n_heads: number of heads in self-attention
  • dropout: dropout rate, default 0.1

TFEncoder.forward(x, mask)

  • x (~torch.FloatTensor): shape (batch_size, max_seq_len, d_model)
  • mask (~torch.ByteTensor): shape (batch_size, max_seq_len)


import torch
import tfencoder

encoder = tfencoder.TFEncoder(6, 512, 2048, 8, dropout=0.1)

x = torch.randn(64, 100, 512)  # (batch_size, max_seq_len, d_model)
mask = torch.randn(64, 100).ge(0)  # a random mask

out = encoder(x, mask)

This package also provides the embedding, positional encoding and scheduled optimizer that are used in transformer as extra functionalities.

TFEmbedding(d_model, n_vocab)

  • d_model: same as TFEncoder
  • n_vocab: vocabulary size


  • x (~torch.LongTensor): shape (batch_size, max_seq_len)

TFPositionalEncoding(d_model, dropout, max_len)

  • d_model: same as TFEncoder
  • dropout: dropout rate
  • max_len: max sequence length


  • x (~torch.FloatTensor): shape (batch_size, max_seq_len, d_model)


import torch
import torch.nn as nn

from tfencoder.utils import TFEmbedding, TFPositionalEncoding

tfembed = TFEmbedding(512, 6)
tfpe = TFPositionalEncoding(512, 0.1, max_len=5)
tfinput = nn.Sequential(tfembed, tfpe)

x = torch.LongTensor([[1,2,3,4,5], [1,2,3,0,0]])
out = tfinput(x)

TFOptimizer(d_model, factor, warmup, optimizer)

  • d_model: equals d_model in TFEncoder
  • factor: scale factor of learning rate
  • warmup: warmup steps
  • optimizer (~torch.optim.Optimzier): e.g. adam optimzier


import torch.optim as optim

from tfencoder import TFEncoder
from tfencoder.utils import TFOptimizer

encoder = TFEncoder(6, 512, 2048, 8, dropout=0.1)
optimizer = TFOptimizer(512, 1, 1000, optim.Adam(encoder.parameters(), lr=0))


loss = ...



Any contributions are welcome!

Project details

Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Files for tfencoder, version 0.0.2
Filename, size File type Python version Upload date Hashes
Filename, size tfencoder-0.0.2-py3-none-any.whl (9.2 kB) File type Wheel Python version py3 Upload date Hashes View hashes
Filename, size tfencoder-0.0.2.tar.gz (5.9 kB) File type Source Python version None Upload date Hashes View hashes

Supported by

Elastic Elastic Search Pingdom Pingdom Monitoring Google Google BigQuery Sentry Sentry Error logging AWS AWS Cloud computing DataDog DataDog Monitoring Fastly Fastly CDN SignalFx SignalFx Supporter DigiCert DigiCert EV certificate StatusPage StatusPage Status page