a-One-Fan/DL-Art-School
1
0
Fork 0
forked from mrq/DL-Art-School
Code Pull Requests Activity

Rework minicoder to always checkpoint

Browse Source
This commit is contained in:
James Betker 2022-03-01 14:09:18 -07:00
parent db0c3340ac
commit 45ab444c04
1 changed files with 6 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
codes/models/gpt_voice/mini_encoder.py
View File

@ -7,7 +7,7 @@ from models.diffusion.unet_diffusion import Downsample, AttentionBlock, QKVAtten
# Combined resnet & full-attention encoder for converting an audio clip into an embedding. # Combined resnet & full-attention encoder for converting an audio clip into an embedding.
from trainer.networks import register_model from trainer.networks import register_model
from utils.util import checkpoint, opt_get from utils.util import checkpoint, opt_get, sequential_checkpoint
class ResBlock(nn.Module): class ResBlock(nn.Module):
@ -100,14 +100,14 @@ class AudioMiniEncoder(nn.Module):
num_attn_heads=4, num_attn_heads=4,
dropout=0, dropout=0,
downsample_factor=2, downsample_factor=2,
kernel_size=3, kernel_size=3):
do_checkpointing=False):
super().__init__() super().__init__()
self.init = nn.Sequential( self.init = nn.Sequential(
conv_nd(1, spec_dim, base_channels, 3, padding=1) conv_nd(1, spec_dim, base_channels, 3, padding=1)
) )
ch = base_channels ch = base_channels
res = [] res = []
self.layers = depth
for l in range(depth): for l in range(depth):
for r in range(resnet_blocks): for r in range(resnet_blocks):
res.append(ResBlock(ch, dropout, dims=1, do_checkpoint=False, kernel_size=kernel_size)) res.append(ResBlock(ch, dropout, dims=1, do_checkpoint=False, kernel_size=kernel_size))
@ -124,16 +124,13 @@ class AudioMiniEncoder(nn.Module):
attn.append(AttentionBlock(embedding_dim, num_attn_heads, do_checkpoint=False)) attn.append(AttentionBlock(embedding_dim, num_attn_heads, do_checkpoint=False))
self.attn = nn.Sequential(*attn) self.attn = nn.Sequential(*attn)
self.dim = embedding_dim self.dim = embedding_dim
self.do_checkpointing = do_checkpointing
def forward(self, x): def forward(self, x):
h = self.init(x) h = self.init(x)
h = self.res(h) h = sequential_checkpoint(self.res, self.layers, h)
h = self.final(h) h = self.final(h)
if self.do_checkpointing: for blk in self.attn:
h = checkpoint(self.attn, h) h = checkpoint(blk, h)
else:
h = self.attn(h)
return h[:, :, 0] return h[:, :, 0]