def __init__(self,
dec_inp_size,
dec_out_size,
d_latent,
N=6,
d_model=512,
d_ff=2048,
h=8,
dropout=0.1,
device='cpu',
d_map_latent=8):
super(DecoderY, self).__init__()
self.dec_out_size = dec_out_size
self.d_model = d_model
self.device = device
self.trg_embed = nn.Sequential(
LinearEmbedding(dec_inp_size, d_model - d_map_latent),
PositionalEncoding(d_model - d_map_latent, dropout))
self.decoder = Decoder(
DecoderLayer(
d_model, MultiHeadAttention(h, d_model),
MultiHeadAttention(h, d_model),
ConcatPointerwiseFeedforward(d_model, d_latent, d_ff, dropout),
dropout), N)
self.fc = nn.Linear(d_model, dec_out_size * 2)
self.init_weights(self.decoder.parameters())
self.init_weights(self.fc.parameters())
self.map_encoder = load_map_encoder(device)
def __init__(self,
enc_inp_size,
dec_inp_size,
dec_out_size,
N=6,
d_model=512,
d_ff=2048,
heads=8,
dropout=0.1,
mean=[0, 0],
std=[0, 0]):
super(IndividualTF, self).__init__()
"Helper: Construct a model from hyperparameters."
c = copy.deepcopy
attn = MultiHeadAttention(heads, d_model)
ff = PointerwiseFeedforward(d_model, d_ff, dropout)
position = PositionalEncoding(d_model, dropout)
self.mean = np.array(mean)
self.std = np.array(std)
self.model = EncoderDecoder(
Encoder(EncoderLayer(d_model, c(attn), c(ff), dropout), N),
Decoder(DecoderLayer(d_model, c(attn), c(attn), c(ff), dropout),
N),
nn.Sequential(LinearEmbedding(enc_inp_size, d_model), c(position)),
nn.Sequential(LinearEmbedding(dec_inp_size, d_model), c(position)),
Generator(d_model, dec_out_size))
# This was important from their code.
# Initialize parameters with Glorot / fan_avg.
for p in self.model.parameters():
if p.dim() > 1:
nn.init.xavier_uniform_(p)
def __init__(self,
src_len,
tgt_len,
enc_inp_size,
dec_inp_size,
dec_out_size,
N=6,
d_model=512,
d_ff=2048,
h=8,
dropout=0.1,
device='cpu'):
super(Generator, self).__init__()
self.device = device
self.src_len = src_len
self.tgt_len = tgt_len
self.dec_inp_size = dec_inp_size
c = copy.deepcopy
attn = MultiHeadAttention(h, d_model)
ff = PointerwiseFeedforward(d_model, d_ff, dropout)
position = PositionalEncoding(d_model, dropout)
self.generator = EncoderDecoder(
Encoder(EncoderLayer(d_model, c(attn), c(ff), dropout), N),
Decoder(DecoderLayer(d_model, c(attn), c(attn), c(ff), dropout),
N),
nn.Sequential(LinearEmbedding(enc_inp_size, d_model), c(position)),
nn.Sequential(LinearEmbedding(dec_inp_size, d_model), c(position)),
TFHeadGenerator(d_model, dec_out_size))
# This was important from their code.
# Initialize parameters with Glorot / fan_avg.
for p in self.generator.parameters():
if p.dim() > 1:
nn.init.xavier_uniform_(p)
def __init__(self,
disc_inp_size,
disc_seq_len,
N=6,
d_model=512,
d_ff=2048,
h=8,
dropout=0.1,
device='cpu'):
super(Critic, self).__init__()
self.device = device
c = copy.deepcopy
attn = MultiHeadAttention(h, d_model)
ff = PointerwiseFeedforward(d_model, d_ff, dropout)
position = PositionalEncoding(d_model, dropout)
self.critic = nn.ModuleDict({
'src_embed':
nn.Sequential(LinearEmbedding(disc_inp_size, d_model),
c(position)),
'encoder':
Encoder(EncoderLayer(d_model, c(attn), c(ff), dropout), N),
'disc_head':
nn.Sequential(nn.Flatten(), nn.Linear(d_model * disc_seq_len, 1)),
})
for p in self.critic.parameters():
if p.dim() > 1:
nn.init.xavier_uniform_(p)