Skip to content

Instantly share code, notes, and snippets.

@MasanoriYamada
Last active January 23, 2020 11:54
Show Gist options
  • Save MasanoriYamada/d1d8ca884d200e73cca66a4387c7470a to your computer and use it in GitHub Desktop.
Save MasanoriYamada/d1d8ca884d200e73cca66a4387c7470a to your computer and use it in GitHub Desktop.
import torch
def get_batch_jacobian(net, x, to):
# noutputs: total output dim (e.g. net(x).shape(b,1,4,4) noutputs=1*4*4
# b: batch
# i: in_dim
# o: out_dim
# ti: total input dim
# to: total output dim
x_batch = x.shape[0]
x_shape = x.shape[1:]
x = x.unsqueeze(1) # b, 1 ,i
x = x.repeat(1, to, *(1,)*len(x.shape[2:])) # b * to,i copy to o dim
x.requires_grad_(True)
tmp_shape = x.shape
y = net(x.reshape(-1, *tmp_shape[2:])) # x.shape = b*to,i y.shape = b*to,to
y_shape = y.shape[1:] # y.shape = b*to,to
y = y.reshape(x_batch, to, to) # y.shape = b,to,to
input_val = torch.eye(to).reshape(1, to, to).repeat(x_batch, 1, 1) # input_val.shape = b,to,to value is (eye)
y.backward(input_val) # y.shape = b,to,to
return x.grad.reshape(x_batch, *y_shape, *x_shape).data # x.shape = b,o,i
class CNNNet(torch.nn.Module):
def __init__(self):
super(CNNNet, self).__init__()
self.cnn = torch.nn.Conv2d(1, 3, 5)
self.fc1 = torch.nn.Linear(3, 4)
def forward(self, x):
print('x: {}'.format(x.shape))
x = torch.nn.functional.relu(self.cnn(x))
print('co: {}'.format(x.shape))
#x = x.reshape(x.shape[0], -1)
#x = torch.nn.functional.relu(self.fc1(x))
#print('li: {}'.format(x.shape))
return x
cnet = CNNNet()
batch = 10
x = torch.randn(batch,1,5,5)
y = cnet(x)
ret = get_batch_jacobian(cnet, x, 3) # y.shape=10,3,1.1
print(ret.shape)
@RylanSchaeffer
Copy link

RylanSchaeffer commented Jan 18, 2020

I have a question. What would the consequence be of not repeating input x before the forward pass, but repeating the output y after the forward pass? What would the resulting .grad field contain?

@MasanoriYamada
Copy link
Author

MasanoriYamada commented Jan 23, 2020

@RylanSchaeffer

I believe that in the comment, y.shape should be b*to, to. Just for anyone else that comes along after me!

Thank you!
Could you show me your complete code for reproducing your error?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment