From 5fbf447e2d8e336683c96b6a074501c72e8c59e8 Mon Sep 17 00:00:00 2001 From: Yin Li Date: Wed, 4 Mar 2020 16:21:00 -0500 Subject: [PATCH] Fix spectral norm interference with gradient visualization spectral norm introduces weight_orig, so that not all weight parameters ends with ".weight" --- map2map/train.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/map2map/train.py b/map2map/train.py index 051361a..ce61f79 100644 --- a/map2map/train.py +++ b/map2map/train.py @@ -364,7 +364,7 @@ def train(epoch, loader, model, criterion, optimizer, scheduler, # gradients of the weights of the first and the last layer grads = list(p.grad for n, p in model.named_parameters() - if n.endswith('weight')) + if '.weight' in n) grads = [grads[0], grads[-1]] grads = [g.detach().norm().item() for g in grads] logger.add_scalars('grad', { @@ -373,7 +373,7 @@ def train(epoch, loader, model, criterion, optimizer, scheduler, }, global_step=batch) if args.adv and epoch >= args.adv_start: grads = list(p.grad for n, p in adv_model.named_parameters() - if n.endswith('weight')) + if '.weight' in n) grads = [grads[0], grads[-1]] grads = [g.detach().norm().item() for g in grads] logger.add_scalars('grad/adv', {