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

Make vqvae3_hard more configurable

Browse Source
This commit is contained in:
James Betker 2021-02-04 09:03:22 -07:00
parent b980028ca8
commit 7070142805
3 changed files with 42 additions and 30 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
codes/models/switched_conv/switched_conv_hard_routing.py
View File

@ -146,14 +146,16 @@ class SwitchNorm(nn.Module):
class HardRoutingGate(nn.Module): class HardRoutingGate(nn.Module):
def __init__(self, breadth): def __init__(self, breadth, hard_en=True):
super().__init__() super().__init__()
self.norm = SwitchNorm(breadth, accumulator_size=256) self.norm = SwitchNorm(breadth, accumulator_size=256)
self.hard_en = hard_en
def forward(self, x): def forward(self, x):
soft = self.norm(nn.functional.softmax(x, dim=1)) soft = self.norm(nn.functional.softmax(x, dim=1))
hard = RouteTop1.apply(soft) # This variant can route gradients downstream. if self.hard_en:
return hard return RouteTop1.apply(soft)
return soft
class SwitchedConvHardRouting(nn.Module): class SwitchedConvHardRouting(nn.Module):
@ -167,7 +169,8 @@ class SwitchedConvHardRouting(nn.Module):
dropout_rate=0.0, dropout_rate=0.0,
include_coupler: bool = False, # A 'coupler' is a latent converter which can make any bxcxhxw tensor a compatible switchedconv selector by performing a linear 1x1 conv, softmax and interpolate. include_coupler: bool = False, # A 'coupler' is a latent converter which can make any bxcxhxw tensor a compatible switchedconv selector by performing a linear 1x1 conv, softmax and interpolate.
coupler_mode: str = 'standard', coupler_mode: str = 'standard',
coupler_dim_in: int = 0): coupler_dim_in: int = 0,
hard_en=True): # A test switch that, when used in 'emulation mode' (where all convs are calculated using torch functions) computes soft-attention instead of hard-attention.
super().__init__() super().__init__()
self.in_channels = in_c self.in_channels = in_c
self.out_channels = out_c self.out_channels = out_c
@ -190,7 +193,7 @@ class SwitchedConvHardRouting(nn.Module):
Conv2d(breadth, breadth, 1, stride=self.stride)) Conv2d(breadth, breadth, 1, stride=self.stride))
else: else:
self.coupler = None self.coupler = None
self.gate = HardRoutingGate(breadth) self.gate = HardRoutingGate(breadth, hard_en=hard_en)
self.weight = nn.Parameter(torch.empty(out_c, in_c, breadth, kernel_sz, kernel_sz)) self.weight = nn.Parameter(torch.empty(out_c, in_c, breadth, kernel_sz, kernel_sz))
if bias: if bias:

57
codes/models/vqvae/vqvae_3_hardswitch.py
View File

@ -15,10 +15,10 @@ from utils.util import checkpoint, opt_get
# Upsamples and blurs (similar to StyleGAN). Replaces ConvTranspose2D from the original paper. # Upsamples and blurs (similar to StyleGAN). Replaces ConvTranspose2D from the original paper.
class UpsampleConv(nn.Module): class UpsampleConv(nn.Module):
def __init__(self, in_filters, out_filters, kernel_size, padding): def __init__(self, in_filters, out_filters, kernel_size, padding, cfg):
super().__init__() super().__init__()
self.conv = SwitchedConvHardRouting(in_filters, out_filters, kernel_size, breadth=16, include_coupler=True, self.conv = SwitchedConvHardRouting(in_filters, out_filters, kernel_size, breadth=cfg['breadth'], include_coupler=True,
coupler_mode='standard', coupler_dim_in=in_filters, dropout_rate=0.4) coupler_mode=cfg['mode'], coupler_dim_in=in_filters, dropout_rate=cfg['dropout'], hard_en=cfg['hard_enabled'])
def forward(self, x): def forward(self, x):
up = torch.nn.functional.interpolate(x, scale_factor=2) up = torch.nn.functional.interpolate(x, scale_factor=2)
@ -26,26 +26,26 @@ class UpsampleConv(nn.Module):
class Encoder(nn.Module): class Encoder(nn.Module):
def __init__(self, in_channel, channel, n_res_block, n_res_channel, stride): def __init__(self, in_channel, channel, n_res_block, n_res_channel, stride, cfg):
super().__init__() super().__init__()
if stride == 4: if stride == 4:
blocks = [ blocks = [
nn.Conv2d(in_channel, channel // 2, 5, stride=2, padding=2), nn.Conv2d(in_channel, channel // 2, 5, stride=2, padding=2),
nn.LeakyReLU(inplace=True), nn.LeakyReLU(inplace=True),
SwitchedConvHardRouting(channel // 2, channel, 5, breadth=16, stride=2, include_coupler=True, SwitchedConvHardRouting(channel // 2, channel, 5, breadth=cfg['breadth'], stride=2, include_coupler=True,
coupler_mode='standard', coupler_dim_in=channel // 2, dropout_rate=0.4), coupler_mode=cfg['mode'], coupler_dim_in=channel // 2, dropout_rate=cfg['dropout'], hard_en=cfg['hard_enabled']),
nn.LeakyReLU(inplace=True), nn.LeakyReLU(inplace=True),
SwitchedConvHardRouting(channel, channel, 3, breadth=16, include_coupler=True, coupler_mode='standard', SwitchedConvHardRouting(channel, channel, 3, breadth=cfg['breadth'], include_coupler=True, coupler_mode=cfg['mode'],
coupler_dim_in=channel, dropout_rate=0.4), coupler_dim_in=channel, dropout_rate=cfg['dropout'], hard_en=cfg['hard_enabled']),
] ]
elif stride == 2: elif stride == 2:
blocks = [ blocks = [
nn.Conv2d(in_channel, channel // 2, 5, stride=2, padding=2), nn.Conv2d(in_channel, channel // 2, 5, stride=2, padding=2),
nn.LeakyReLU(inplace=True), nn.LeakyReLU(inplace=True),
SwitchedConvHardRouting(channel // 2, channel, 3, breadth=16, include_coupler=True, coupler_mode='standard', SwitchedConvHardRouting(channel // 2, channel, 3, breadth=cfg['breadth'], include_coupler=True, coupler_mode=cfg['mode'],
coupler_dim_in=channel // 2, dropout_rate=0.4), coupler_dim_in=channel // 2, dropout_rate=cfg['dropout'], hard_en=cfg['hard_enabled']),
] ]
for i in range(n_res_block): for i in range(n_res_block):
@ -61,12 +61,12 @@ class Encoder(nn.Module):
class Decoder(nn.Module): class Decoder(nn.Module):
def __init__( def __init__(
self, in_channel, out_channel, channel, n_res_block, n_res_channel, stride self, in_channel, out_channel, channel, n_res_block, n_res_channel, stride, cfg
): ):
super().__init__() super().__init__()
blocks = [SwitchedConvHardRouting(in_channel, channel, 3, breadth=16, include_coupler=True, coupler_mode='standard', blocks = [SwitchedConvHardRouting(in_channel, channel, 3, breadth=cfg['breadth'], include_coupler=True, coupler_mode=cfg['mode'],
coupler_dim_in=in_channel, dropout_rate=0.4)] coupler_dim_in=in_channel, dropout_rate=cfg['dropout'], hard_en=cfg['hard_enabled'])]
for i in range(n_res_block): for i in range(n_res_block):
blocks.append(ResBlock(channel, n_res_channel)) blocks.append(ResBlock(channel, n_res_channel))
@ -76,17 +76,17 @@ class Decoder(nn.Module):
if stride == 4: if stride == 4:
blocks.extend( blocks.extend(
[ [
UpsampleConv(channel, channel // 2, 5, padding=2), UpsampleConv(channel, channel // 2, 5, padding=2, cfg=cfg),
nn.LeakyReLU(inplace=True), nn.LeakyReLU(inplace=True),
UpsampleConv( UpsampleConv(
channel // 2, out_channel, 5, padding=2 channel // 2, out_channel, 5, padding=2, cfg=cfg
), ),
] ]
) )
elif stride == 2: elif stride == 2:
blocks.append( blocks.append(
UpsampleConv(channel, out_channel, 5, padding=2) UpsampleConv(channel, out_channel, 5, padding=2, cfg=cfg)
) )
self.blocks = nn.Sequential(*blocks) self.blocks = nn.Sequential(*blocks)
@ -105,22 +105,24 @@ class VQVAE3HardSwitch(nn.Module):
codebook_dim=64, codebook_dim=64,
codebook_size=512, codebook_size=512,
decay=0.99, decay=0.99,
cfg={'mode':'standard', 'breadth':16, 'hard_enabled': True, 'dropout': 0.4}
): ):
super().__init__() super().__init__()
self.cfg = cfg
self.initial_conv = nn.Sequential(*[nn.Conv2d(in_channel, 32, 3, padding=1), self.initial_conv = nn.Sequential(*[nn.Conv2d(in_channel, 32, 3, padding=1),
nn.LeakyReLU(inplace=True)]) nn.LeakyReLU(inplace=True)])
self.enc_b = Encoder(32, channel, n_res_block, n_res_channel, stride=4) self.enc_b = Encoder(32, channel, n_res_block, n_res_channel, stride=4, cfg=cfg)
self.enc_t = Encoder(channel, channel, n_res_block, n_res_channel, stride=2) self.enc_t = Encoder(channel, channel, n_res_block, n_res_channel, stride=2, cfg=cfg)
self.quantize_conv_t = nn.Conv2d(channel, codebook_dim, 1) self.quantize_conv_t = nn.Conv2d(channel, codebook_dim, 1)
self.quantize_t = Quantize(codebook_dim, codebook_size) self.quantize_t = Quantize(codebook_dim, codebook_size)
self.dec_t = Decoder( self.dec_t = Decoder(
codebook_dim, codebook_dim, channel, n_res_block, n_res_channel, stride=2 codebook_dim, codebook_dim, channel, n_res_block, n_res_channel, stride=2, cfg=cfg
) )
self.quantize_conv_b = nn.Conv2d(codebook_dim + channel, codebook_dim, 1) self.quantize_conv_b = nn.Conv2d(codebook_dim + channel, codebook_dim, 1)
self.quantize_b = Quantize(codebook_dim, codebook_size) self.quantize_b = Quantize(codebook_dim, codebook_size)
self.upsample_t = UpsampleConv( self.upsample_t = UpsampleConv(
codebook_dim, codebook_dim, 5, padding=2 codebook_dim, codebook_dim, 5, padding=2, cfg=cfg
) )
self.dec = Decoder( self.dec = Decoder(
codebook_dim + codebook_dim, codebook_dim + codebook_dim,
@ -129,6 +131,7 @@ class VQVAE3HardSwitch(nn.Module):
n_res_block, n_res_block,
n_res_channel, n_res_channel,
stride=4, stride=4,
cfg=cfg
) )
self.final_conv = nn.Conv2d(32, in_channel, 3, padding=1) self.final_conv = nn.Conv2d(32, in_channel, 3, padding=1)
@ -211,7 +214,7 @@ def convert_weights(weights_file):
from models.vqvae.vqvae_3 import VQVAE3 from models.vqvae.vqvae_3 import VQVAE3
std_model = VQVAE3() std_model = VQVAE3()
std_model.load_state_dict(sd) std_model.load_state_dict(sd)
nsd = convert_conv_net_state_dict_to_switched_conv(std_model, 16, ['quantize_conv_t', 'quantize_conv_b', nsd = convert_conv_net_state_dict_to_switched_conv(std_model, 8, ['quantize_conv_t', 'quantize_conv_b',
'enc_b.blocks.0', 'enc_t.blocks.0', 'enc_b.blocks.0', 'enc_t.blocks.0',
'conv.1', 'conv.3', 'initial_conv', 'final_conv']) 'conv.1', 'conv.3', 'initial_conv', 'final_conv'])
torch.save(nsd, "converted.pth") torch.save(nsd, "converted.pth")
@ -224,7 +227,13 @@ def register_vqvae3_hard_switch(opt_net, opt):
def performance_test(): def performance_test():
net = VQVAE3HardSwitch().to('cuda') cfg = {
'mode': 'lambda',
'breadth': 8,
'hard_enabled': False,
'dropout': 0
}
net = VQVAE3HardSwitch(cfg=cfg).to('cuda')
loss = nn.L1Loss() loss = nn.L1Loss()
opt = torch.optim.Adam(net.parameters(), lr=1e-4) opt = torch.optim.Adam(net.parameters(), lr=1e-4)
started = time() started = time()
@ -241,5 +250,5 @@ def performance_test():
if __name__ == '__main__': if __name__ == '__main__':
#v = VQVAE3HardSwitch() #v = VQVAE3HardSwitch()
#print(v(torch.randn(1,3,128,128))[0].shape) #print(v(torch.randn(1,3,128,128))[0].shape)
#convert_weights("../../../experiments/test_vqvae3.pth") convert_weights("../../../experiments/test_vqvae3.pth")
performance_test() #performance_test()

2
codes/train.py
View File

@ -295,7 +295,7 @@ class Trainer:
if __name__ == '__main__': if __name__ == '__main__':
parser = argparse.ArgumentParser() parser = argparse.ArgumentParser()
parser.add_argument('-opt', type=str, help='Path to option YAML file.', default='../options/train_imgset_tiled_nvqvae_stage1_lambda.yml') parser.add_argument('-opt', type=str, help='Path to option YAML file.', default='../options/train_imgset_vqvae3_stage1.yml')
parser.add_argument('--launcher', choices=['none', 'pytorch'], default='none', help='job launcher') parser.add_argument('--launcher', choices=['none', 'pytorch'], default='none', help='job launcher')
parser.add_argument('--local_rank', type=int, default=0) parser.add_argument('--local_rank', type=int, default=0)
args = parser.parse_args() args = parser.parse_args()