File size: 1,350 Bytes
d639276 7494603 d639276 c646511 d639276 7494603 d639276 7494603 d639276 7494603 d639276 c646511 d639276 7494603 d639276 7494603 d639276 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 |
import torch
import torch.nn as nn
from nets.decoder import Decoder
from nets.projections import Projections
from nets.encoder import Encoder
class Model(nn.Module):
def __init__(self, input_size, embedding_size,
decoder_input_size,
num_heads=8, num_layers=4, ff_hidden=250, *args, **kwargs):
super().__init__()
self.embedding_size = embedding_size
# ----------- Encoder -----------
self.encoder = Encoder(
n_heads=num_heads,
embed_dim=embedding_size,
n_layers=num_layers,
feed_forward_hidden=ff_hidden,
node_dim=input_size
)
# ----------- Decoder -----------
self.decoder = Decoder(
decoder_input_size=decoder_input_size,
embedding_size=embedding_size,
num_heads=num_heads
)
# ----------- Attention Projections -----------
self.projections = Projections(
n_heads=num_heads,
embed_dim=embedding_size
)
# ----------- Fleet Attention Encoder (Optional) -----------
self.fleet_attention = Encoder(
n_heads=num_heads,
embed_dim=embedding_size,
n_layers=1,
feed_forward_hidden=ff_hidden,
node_dim=embedding_size + 1
)
|