diff options
author | Marcin Chrzanowski <m@m-chrzan.xyz> | 2021-05-27 19:54:24 +0200 |
---|---|---|
committer | Marcin Chrzanowski <m@m-chrzan.xyz> | 2021-05-27 19:54:24 +0200 |
commit | bb6a6004bfb5398b5cdf8f8a973b47e5659aec79 (patch) | |
tree | f770a9cdb856d25e731289c86583ac3c71937f3a /model | |
parent | 22cd84cf78a114d75a93f75f66f6ea61c02c94ef (diff) |
Fix indents
Diffstat (limited to 'model')
-rw-r--r-- | model/encoder.py | 90 |
1 files changed, 49 insertions, 41 deletions
diff --git a/model/encoder.py b/model/encoder.py index 63a5149..85b3141 100644 --- a/model/encoder.py +++ b/model/encoder.py @@ -4,45 +4,53 @@ from torch import nn from util.util import get_positional_encoding from model.encoder_layer import EncoderLayer +def log(string, verbose): + if verbose: + print(string) + class EncoderModel(nn.Module): - def __init__(self, input_dim, hidden_dim, d_ff, output_dim, n_layers, - num_heads, use_attention=True, use_feedforward=True, - use_positional=True, device='cpu'): - super(EncoderModel, self).__init__() - self._device = device - self._use_positional = use_positional - self.embedding_layer = nn.Embedding(input_dim, hidden_dim) - self.layers = nn.ModuleList([ - EncoderLayer(hidden_dim, d_ff, num_heads, use_attention, - use_feedforward) for i in range(n_layers)]) - self.output_layer = nn.Linear(hidden_dim, output_dim) - - def forward(self, x, return_att_weights=False): - # x shape: (seqlen, batch) - hidden = self.embedding_layer(x) - # hidden shape: (seqlen, batch, hiddendim) - - if self._use_positional: - positional_encoding = get_positional_encoding( - n_positions=hidden.shape[0], - n_dimensions=hidden.shape[-1], - device=self._device - ) - # reshaping to (seqlen, 1, hiddendim) - positional_encoding = torch.reshape( - positional_encoding, - (hidden.shape[0], 1, hidden.shape[-1]) - ) - hidden = hidden + positional_encoding - - list_att_weights = [] - for layer in self.layers: - hidden, att_weights = layer(hidden) - list_att_weights.append(att_weights) - - result = self.output_layer(hidden) - - if return_att_weights: - return result, list_att_weights - else: - return result + def __init__(self, input_dim, hidden_dim, d_ff, output_dim, n_layers, + num_heads, use_attention=True, use_feedforward=True, + use_positional=True, device='cpu'): + super(EncoderModel, self).__init__() + self._device = device + self._use_positional = use_positional + self.embedding_layer = nn.Embedding(input_dim, hidden_dim) + self.layers = nn.ModuleList([ + EncoderLayer(hidden_dim, d_ff, num_heads, use_attention, + use_feedforward) for i in range(n_layers) + ]) + self.output_layer = nn.Linear(hidden_dim, output_dim) + + def forward(self, x, return_att_weights=False, verbose=False): + log(f'Handling {x}', verbose) + # x shape: (seqlen, batch) + hidden = self.embedding_layer(x) + # hidden shape: (seqlen, batch, hiddendim) + + if self._use_positional: + positional_encoding = get_positional_encoding( + n_positions=hidden.shape[0], + n_dimensions=hidden.shape[-1], + device=self._device + ) + # reshaping to (seqlen, 1, hiddendim) + positional_encoding = torch.reshape( + positional_encoding, + (hidden.shape[0], 1, hidden.shape[-1]) + ) + hidden = hidden + positional_encoding + + list_att_weights = [] + for layer in self.layers: + hidden, att_weights = layer(hidden) + list_att_weights.append(att_weights) + + result = self.output_layer(hidden) + + log('Result: {result}', verbose) + + if return_att_weights: + return result, list_att_weights + else: + return result |