Fix bug from 1963530 to show last instead of 2nd layer gradients

This commit is contained in:
Yin Li 2020-03-03 20:53:34 -05:00
parent d768aadb36
commit 33fd13e4ac

View File

@ -365,20 +365,20 @@ def train(epoch, loader, model, criterion, optimizer, scheduler,
# gradients of the weights of the first and the last layer # gradients of the weights of the first and the last layer
grads = list(p.grad for n, p in model.named_parameters() grads = list(p.grad for n, p in model.named_parameters()
if n.endswith('weight')) if n.endswith('weight'))
grads = [grads[0], grads[1]] grads = [grads[0], grads[-1]]
grads = [g.detach().norm().item() for g in grads] grads = [g.detach().norm().item() for g in grads]
logger.add_scalars('grad', { logger.add_scalars('grad', {
'first': grads[0], 'first': grads[0],
'last': grads[1], 'last': grads[-1],
}, global_step=batch) }, global_step=batch)
if args.adv and epoch >= args.adv_start: if args.adv and epoch >= args.adv_start:
grads = list(p.grad for n, p in adv_model.named_parameters() grads = list(p.grad for n, p in adv_model.named_parameters()
if n.endswith('weight')) if n.endswith('weight'))
grads = [grads[0], grads[1]] grads = [grads[0], grads[-1]]
grads = [g.detach().norm().item() for g in grads] grads = [g.detach().norm().item() for g in grads]
logger.add_scalars('grad/adv', { logger.add_scalars('grad/adv', {
'first': grads[0], 'first': grads[0],
'last': grads[1], 'last': grads[-1],
}, global_step=batch) }, global_step=batch)
if args.adv and epoch >= args.adv_start: if args.adv and epoch >= args.adv_start: