From b980028ca802d0239d51a9e0f43510e18c62a22f Mon Sep 17 00:00:00 2001 From: James Betker Date: Wed, 3 Feb 2021 14:12:24 -0700 Subject: [PATCH] Add get_debug_values for vqvae_3_hardswitch --- codes/models/vqvae/vqvae_3_hardswitch.py | 38 ++++++++++++++++++++++-- 1 file changed, 36 insertions(+), 2 deletions(-) diff --git a/codes/models/vqvae/vqvae_3_hardswitch.py b/codes/models/vqvae/vqvae_3_hardswitch.py index 1222a5ba..958eda98 100644 --- a/codes/models/vqvae/vqvae_3_hardswitch.py +++ b/codes/models/vqvae/vqvae_3_hardswitch.py @@ -1,6 +1,10 @@ +import os +from time import time + import torch import torchvision from torch import nn +from tqdm import tqdm from models.switched_conv.switched_conv_hard_routing import SwitchedConvHardRouting, \ convert_conv_net_state_dict_to_switched_conv @@ -146,7 +150,21 @@ class VQVAE3HardSwitch(nn.Module): def visual_dbg(self, step, path): convs = [self.dec.blocks[-1].conv, self.dec_t.blocks[-1].conv, self.enc_b.blocks[-4], self.enc_t.blocks[-4]] for i, c in enumerate(convs): - self.save_attention_to_image_rgb(os.path.join(path, "%i_selector_%i.png" % (step, i+1)), c.last_select, self.breadth) + self.save_attention_to_image_rgb(os.path.join(path, "%i_selector_%i.png" % (step, i+1)), c.last_select, 16) + + def get_debug_values(self, step, __): + switched_convs = [('enc_b_blk2', self.enc_b.blocks[2]), + ('enc_b_blk4', self.enc_b.blocks[4]), + ('enc_t_blk2', self.enc_t.blocks[2]), + ('dec_t_blk0', self.dec_t.blocks[0]), + ('dec_t_blk-1', self.dec_t.blocks[-1].conv), + ('dec_blk0', self.dec.blocks[0]), + ('dec_blk-1', self.dec.blocks[-1].conv), + ('dec_blk-3', self.dec.blocks[-3].conv)] + logs = {} + for name, swc in switched_convs: + logs[f'{name}_histogram_switch_usage'] = swc.latest_masks + return logs def encode(self, input): fea = self.initial_conv(input) @@ -205,7 +223,23 @@ def register_vqvae3_hard_switch(opt_net, opt): return VQVAE3HardSwitch(**kw) +def performance_test(): + net = VQVAE3HardSwitch().to('cuda') + loss = nn.L1Loss() + opt = torch.optim.Adam(net.parameters(), lr=1e-4) + started = time() + for j in tqdm(range(10)): + inp = torch.rand((8, 3, 256, 256), device='cuda') + res = net(inp)[0] + l = loss(res, inp) + l.backward() + opt.step() + net.zero_grad() + print("Elapsed: ", (time()-started)) + + if __name__ == '__main__': #v = VQVAE3HardSwitch() #print(v(torch.randn(1,3,128,128))[0].shape) - convert_weights("../../../experiments/test_vqvae3.pth") \ No newline at end of file + #convert_weights("../../../experiments/test_vqvae3.pth") + performance_test()