diff options
| -rw-r--r-- | sdil/core.py | 35 |
1 files changed, 22 insertions, 13 deletions
diff --git a/sdil/core.py b/sdil/core.py index 66f2d0a..3cd8a97 100644 --- a/sdil/core.py +++ b/sdil/core.py @@ -278,19 +278,28 @@ def simultaneous_node_perturbation_targets(net, x, y, sigma=1e-2, with torch.no_grad(): base = net.forward(x) h = base["h"] - qs = [torch.zeros_like(hl) for hl in h[1:-1]] - L0 = loss_ce(h[-1], y) - for _ in range(n_dirs): - xis = [torch.empty_like(q).bernoulli_(0.5).mul_(2).sub_(1) for q in qs] - Lplus = _loss_with_hidden_noise(net, x, y, xis, sigma) - if antithetic: - Lminus = _loss_with_hidden_noise(net, x, y, xis, -sigma) - coeff = (Lplus - Lminus) / (2.0 * sigma) - else: - coeff = (Lplus - L0) / sigma - for l, xi in enumerate(xis): - qs[l].add_(-coeff.unsqueeze(1) * xi) - return [q / n_dirs for q in qs] + batch = x.shape[0] + xis = [torch.empty((n_dirs,) + hl.shape, device=hl.device, dtype=hl.dtype) + .bernoulli_(0.5).mul_(2).sub_(1) for hl in h[1:-1]] + + # Fold direction and antithetic sign into the batch dimension. This is + # mathematically identical to 2*n_dirs small forwards but turns dozens + # of tiny Pascal-GPU launches into one efficient matrix multiplication + # per layer. + copies = 2 * n_dirs if antithetic else n_dirs + x_batch = x.repeat(copies, 1) + y_batch = y.repeat(copies) + if antithetic: + noise_batch = [torch.cat((xi, -xi), dim=0).flatten(0, 1) for xi in xis] + else: + noise_batch = [xi.flatten(0, 1) for xi in xis] + losses = _loss_with_hidden_noise(net, x_batch, y_batch, noise_batch, sigma) + losses = losses.view(copies, batch) + if antithetic: + coeff = (losses[:n_dirs] - losses[n_dirs:]) / (2.0 * sigma) + else: + coeff = (losses - loss_ce(h[-1], y).unsqueeze(0)) / sigma + return [-(coeff.unsqueeze(2) * xi).mean(0) for xi in xis] def _loss_with_hidden_noise(net, x, y, noises, scale): |