I am trying to calculate the mutual information between the hidden layers’ output and input and output using the following code:
def InfoNCE(X, Y, batch_size=256, num_epochs=200, dev=torch.device(“cpu”), model=None, rg=True):
A = torch.tensor([float(batch_size)] * batch_size).reshape(batch_size, 1)#.cuda()
if not model:
model = nn.Sequential(
nn.Linear(X.shape[1]+Y.shape[1], 16),
nn.ReLU(),
nn.Linear(16, 8),
nn.ReLU(),
nn.Linear(8, 1),
)
# Move data to device
X = X.to(dev)
Y = Y.to(dev) + torch.randn_like(Y) * 1e-4
model = model.to(dev)
opt = optim.SGD(model.parameters(), lr=0.03, momentum=0.9)
td = TensorDataset(X, Y)
result = []
for epoch in range(num_epochs):
for x, y in DataLoader(td, batch_size, shuffle=True, drop_last=True):
opt.zero_grad()
top = model(torch.cat([x, y], 1)).flatten()
xiyj = torch.cat([x.repeat_interleave(batch_size,dim=0),y.repeat(batch_size,1)], 1)
bottom = torch.logsumexp(model(xiyj).reshape(batch_size,batch_size), 1) - A.log()
loss = -(top - bottom).mean()
result.append(-loss.item())
loss.backward(retain_graph=rg)
opt.step()
r = torch.tensor(result[-20:]).mean()
#plt.plot(result)
print(r)
return r
InfoNCE(dataset.x, layer_2_log[1])
I tried setting retain_graph both true and false after reading some posts, it still gives the runtime error no matter what.
Any help would be appreciated!
and it’s much simpler for us.