Class: Secryst::Transformer

Inherits:
Torch::NN::Module
  • Object
show all
Defined in:
lib/secryst/transformer.rb

Instance Method Summary collapse

Constructor Details

#initialize(d_model: 512, nhead: 8, num_encoder_layers: 6, num_decoder_layers: 6, dim_feedforward: 2048, dropout: 0.1, activation: 'relu', custom_encoder: nil, custom_decoder: nil, input_vocab_size:, target_vocab_size:) ⇒ Transformer

A transformer model. User is able to modify the attributes as needed. The architecture is based on the paper “Attention Is All You Need”. Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. In Advances in Neural Information Processing Systems, pages 6000-6010. Users can build the BERT(arxiv.org/abs/1810.04805) model with corresponding parameters. Args:

d_model: the number of expected features in the encoder/decoder inputs (default=512).
nhead: the number of heads in the multiheadattention models (default=8).
num_encoder_layers: the number of sub-encoder-layers in the encoder (default=6).
num_decoder_layers: the number of sub-decoder-layers in the decoder (default=6).
dim_feedforward: the dimension of the feedforward network model (default=2048).
dropout: the dropout value (default=0.1).
activation: the activation function of encoder/decoder intermediate layer, relu or gelu (default=relu).
custom_encoder: custom encoder (default=nil).
custom_decoder: custom decoder (default=nil).
input_vocab_size: size of vocabulary for input sequence (number of different possible tokens).
target_vocab_size: size of vocabulary for target sequence (number of different possible tokens).
Examples

>>> transformer_model = Transformer.new(nhead: 16, num_encoder_layers: 12) >>> src = Torch.rand((10, 32, 512)) >>> tgt = Torch.rand((20, 32, 512)) >>> out = transformer_model.call(src, tgt)



28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
 
# File 'lib/secryst/transformer.rb', line 28

def initialize(d_model: 512, nhead: 8, num_encoder_layers: 6, num_decoder_layers: 6,
  dim_feedforward: 2048, dropout: 0.1, activation: 'relu', custom_encoder: nil, custom_decoder: nil, input_vocab_size:, target_vocab_size:)

  super()

  if custom_encoder
    @encoder = custom_encoder
  else
    encoder_layers = num_encoder_layers.times.map { TransformerEncoderLayer.new(d_model, nhead, dim_feedforward: dim_feedforward, dropout: dropout, activation: activation) }
    encoder_norm = Torch::NN::LayerNorm.new(d_model)
    @encoder = TransformerEncoder.new(encoder_layers, encoder_norm, d_model, input_vocab_size, dropout)
  end

  if custom_decoder
    @decoder = custom_decoder
  else
    decoder_layers = num_decoder_layers.times.map { TransformerDecoderLayer.new(d_model, nhead, dim_feedforward: dim_feedforward, dropout: dropout, activation: activation) }
    decoder_norm = Torch::NN::LayerNorm.new(d_model)
    @decoder = TransformerDecoder.new(decoder_layers, decoder_norm, d_model, target_vocab_size, dropout)
  end

  @linear = Torch::NN::Linear.new(d_model, target_vocab_size)
  @softmax = Torch::NN::LogSoftmax.new(dim: -1)
  _reset_parameters()

  @d_model = d_model
  @nhead = nhead

end

Instance Method Details

#_reset_parameters ⇒ Object 



111
112
113
114
115
 
# File 'lib/secryst/transformer.rb', line 111

def _reset_parameters
  parameters.each do |p|
    Torch::NN::Init.xavier_uniform!(p) if p.dim > 1
  end
end

#forward(src, tgt, src_mask: nil, tgt_mask: nil, memory_mask: nil, src_key_padding_mask: nil, tgt_key_padding_mask: nil, memory_key_padding_mask: nil) ⇒ Object

Take in and process masked source/target sequences. Args:

src: the sequence to the encoder (required).
tgt: the sequence to the decoder (required).
src_mask: the additive mask for the src sequence (optional).
tgt_mask: the additive mask for the tgt sequence (optional).
memory_mask: the additive mask for the encoder output (optional).
src_key_padding_mask: the ByteTensor mask for src keys per batch (optional).
tgt_key_padding_mask: the ByteTensor mask for tgt keys per batch (optional).
memory_key_padding_mask: the ByteTensor mask for memory keys per batch (optional).

Shape:

- src: :math:`(S, N, E)`.
- tgt: :math:`(T, N, E)`.
- src_mask: :math:`(S, S)`.
- tgt_mask: :math:`(T, T)`.
- memory_mask: :math:`(T, S)`.
- src_key_padding_mask: :math:`(N, S)`.
- tgt_key_padding_mask: :math:`(N, T)`.
- memory_key_padding_mask: :math:`(N, S)`.
Note: [src/tgt/memory]_mask ensures that position i is allowed to attend the unmasked
positions. If a ByteTensor is provided, the non-zero positions are not allowed to attend
while the zero positions will be unchanged. If a BoolTensor is provided, positions with ``true``
are not allowed to attend while ``false`` values will be unchanged. If a FloatTensor
is provided, it will be added to the attention weight.
[src/tgt/memory]_key_padding_mask provides specified elements in the key to be ignored by
the attention. If a ByteTensor is provided, the non-zero positions will be ignored while the zero
positions will be unchanged. If a BoolTensor is provided, the positions with the
value of ``true`` will be ignored while the position with the value of ``false`` will be unchanged.
- output: :math:`(T, N, E)`.
Note: Due to the multi-head attention architecture in the transformer model,
the output sequence length of a transformer is same as the input sequence
(i.e. target) length of the decode.
where S is the source sequence length, T is the target sequence length, N is the
batch size, E is the feature number

Examples:

>>> output = transformer_model.call(src, tgt, src_mask: src_mask, tgt_mask: tgt_mask)



94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
 
# File 'lib/secryst/transformer.rb', line 94

def forward(src, tgt, src_mask: nil, tgt_mask: nil,
          memory_mask: nil, src_key_padding_mask: nil,
          tgt_key_padding_mask: nil, memory_key_padding_mask: nil)
  if src.size(1) != tgt.size(1)
    raise RuntimeError, "the batch number of src and tgt must be equal"
  end

  memory = @encoder.call(src, mask: src_mask, src_key_padding_mask: src_key_padding_mask)
  output = @decoder.call(tgt, memory, tgt_mask: tgt_mask, memory_mask: memory_mask,
                        tgt_key_padding_mask: tgt_key_padding_mask,
                        memory_key_padding_mask: memory_key_padding_mask)
  output = @linear.call(output)
  output = @softmax.call(output)

  return output
end