mrq/DL-Art-School
1
3
Fork 5
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Fix feedforward

Browse Source
This commit is contained in:
James Betker 2021-10-29 17:27:51 -06:00
parent b476516340
commit 95ca88efce
1 changed files with 5 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
codes/models/gpt_voice/lucidrains_gpt.py
View File

@ -6,6 +6,7 @@ import torch.nn.functional as F
from einops import rearrange
# helpers
from models.arch_util import checkpoint
from models.gpt_voice.reversible import ReversibleSequence, SequentialSequence
from utils.util import sequential_checkpoint
@ -208,8 +209,8 @@ class Transformer(nn.Module):
def forward(self, x, return_intermediates=False):
intermediates = []
for attn, ff in self.layers.layers:
x_ff = attn(x)
x = ff(x_ff)
x_ff = x + checkpoint(attn, x)
x = x + ff(x_ff)
if return_intermediates:
intermediates.append((x_ff, x))
if return_intermediates:
@ -228,9 +229,9 @@ class Transformer(nn.Module):
assert(len(prev_intermediates) == self.depth)
new_intermediates = []
for (attn, ff), (int_ff, int_out) in zip(self.layers.layers, prev_intermediates):
x = attn(x, only_last_two_elements=True)
x = x + attn(x, only_last_two_elements=True)
# Note that (x) is now only the last two element in the set. Conjoin it with the int_ff latent to compute the norm.
x_ff = torch.cat([int_ff[:,:-1], x], dim=1)
x = ff(x_ff, only_last_two_elements=True)
x = x + ff(x_ff, only_last_two_elements=True)
new_intermediates.append((x_ff, x))
return x, new_intermediates