wassname/pytorch-ts
1
0
Fork 0
mirror of https://github.com/wassname/pytorch-ts.git synced 2026-07-19 11:27:25 +08:00
Code Issues Packages Projects Releases Wiki Activity

transfomer takes [T, B, F] tensors

Browse Source
This commit is contained in:
Dr. Kashif Rasul
2020-01-28 11:21:06 +01:00
parent 5141ae69e9
commit 1ab2babbe3
1 changed files with 8 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pts/model/transformer/transformer_network.py
+8 -8
View File
@@ -59,7 +59,7 @@ class TransformerNetwork(nn.Module):
self.encoder_input = nn.Linear(input_size, d_model)
self.decoder_input = nn.Linear(input_size, d_model)
# [B, T, d_model] where d_model / num_heads...
# [B, T, d_model] where d_model / num_heads is int
self.transformer = nn.Transformer(
d_model=d_model,
nhead=num_heads,
@@ -276,12 +276,12 @@ class TransformerTrainingNetwork(TransformerNetwork):
# inputs, axis=1, begin=self.context_length, end=None
# )
# pass through encoder
enc_out = self.transformer.encoder(self.encoder_input(enc_input))
# pass through encoder [T, B, b_model]
enc_out = self.transformer.encoder(self.encoder_input(enc_input).permute(1,0,2))
# input to decoder
dec_output = self.transformer.decoder(
self.decoder_input(dec_input),
self.decoder_input(dec_input).permute(1,0,2),
enc_out, # memory
tgt_mask=self.upper_triangular_mask(
self.prediction_length
@@ -289,7 +289,7 @@ class TransformerTrainingNetwork(TransformerNetwork):
)
# compute loss
distr_args = self.proj_dist_args(dec_output)
distr_args = self.proj_dist_args(dec_output.permute(1,0,2))
distr = self.distr_output.distribution(distr_args, scale=scale)
loss = - distr.log_prob(future_target)
@@ -384,10 +384,10 @@ class TransformerPredictionNetwork(TransformerNetwork):
)
dec_output = self.transformer.decoder(
self.decoder_input(dec_input), repeated_enc_out, None
self.decoder_input(dec_input).permute(1,0,2), repeated_enc_out, None
)
distr_args = self.proj_dist_args(dec_output)
distr_args = self.proj_dist_args(dec_output.permute(1,0,2))
# compute likelihood of target given the predicted parameters
distr = self.distr_output.distribution(distr_args, scale=repeated_scale)
@@ -450,7 +450,7 @@ class TransformerPredictionNetwork(TransformerNetwork):
)
# pass through encoder
enc_out = self.transformer.encoder(self.encoder_input(inputs))
enc_out = self.transformer.encoder(self.encoder_input(inputs).permute(1,0,2))
return self.sampling_decoder(
past_target=past_target,