Class TorchMultiplyLayer

numpy_ml/tests/nn_torch_models.py:347–390

class TorchMultiplyLayer(nn.Module):
    def __init__(self, act_fn, **kwargs):
        super(TorchMultiplyLayer, self).__init__()
        self.act_fn = act_fn

    def forward(self, Xs):
        self.Xs = []
        x = Xs[0].copy()
        if not isinstance(x, torch.Tensor):
            x = torchify(x)

        self.prod = x.clone()
        x.retain_grad()
        self.Xs.append(x)

        for i in range(1, len(Xs)):
            x = Xs[i]
            if not isinstance(x, torch.Tensor):
                x = torchify(x)

            x.retain_grad()
            self.Xs.append(x)
            self.prod *= x

        self.prod.retain_grad()
        self.Y = self.act_fn(self.prod)
        self.Y.retain_grad()
        return self.Y

    def extract_grads(self, X):
        self.forward(X)
        self.loss = self.Y.sum()
        self.loss.backward()
        grads = {
            "Xs": X,
            "Prod": self.prod.detach().numpy(),
            "Y": self.Y.detach().numpy(),
            "dLdY": self.Y.grad.numpy(),
            "dLdProd": self.prod.grad.numpy(),
        }
        grads.update(
            {"dLdX{}".format(i + 1): xi.grad.numpy() for i, xi in enumerate(self.Xs)}
        )
        return grads


class TorchSkipConnectionIdentity(nn.Module):