- data = []
- for t in range(y.size(1)):
- for d in torch.randperm(y.size(2))[:8]:
- g = torch.autograd.grad(y[0, t, d], x, retain_graph=True)[0]
- sg = g.pow(2).sum().item()
- # sg = 0
- # for p in model.parameters():
- # g = torch.autograd.grad(y[0, t, d], p, retain_graph=True)[0]
- # sg = sg + g.pow(2).sum().item()
- data.append([t, sg])
-
- data = torch.tensor(data)
- ax.scatter(
- data[:, 0], data[:, 1], s=1, label=label
- ) # , color='gray', label='Input')
+ for n, p in [("input", x)] + list(model.named_parameters()):
+ print(f"Processing {model}.{n}")
+ data = []
+ for t in range(y.size(1)):
+ sg = 0
+ for d in torch.randperm(y.size(2))[:8]:
+ sg += torch.autograd.grad(y[0, t, d], p, retain_graph=True)[0]
+ assert not sg.isinf().any()
+ assert not sg.isnan().any()
+ data.append([t, sg.sum().item()])
+
+ data = torch.tensor(data)
+ # cx, cy = data[:, 0], data[:, 1]
+ cy = data[:, 1].sort().values
+ cx = torch.linspace(0, 1, cy.size(0))
+ ax.plot(
+ cx, cy, label=label + "." + n, linewidth=thickness
+ ) # , color='gray', label='Input')