# This shows the weights for each of the 64 hidden neurons and give an idea how each neuron is activated.
learned_weights = rbm_mnist.W.transpose(0,1).numpy()
plt.show()
fig = plt.figure(3, figsize=(10, 10))
for i in range(25):
sub = fig.add_subplot(5, 5, i+1)
sub.imshow(learned_weights[i, :].reshape((28, 28)), cmap=plt.cm.gray)
plt.title(epochs)
plt.show()
#Lets try reconstructing a random number from this model which has learned 5
idx = 7
img = train_dataset.train_data[idx]
reconstructed_img = img.view(-1).type(torch.FloatTensor)
# _ , reconstructed_img1 = rbm_mnist.to_hidden(reconstructed_img)
# _ , reconstructed_img2 = rbm_mnist.to_visible(reconstructed_img)
_, reconstructed_img3 = rbm_mnist.reconstruct(reconstructed_img, 1)
reconstructed_img = reconstructed_img3.view((28, 28))
print("The original number: {}".format(train_dataset.train_labels[idx]))
plt.imshow(img, cmap='gray')
plt.show()
print("The reconstructed image")
plt.imshow(reconstructed_img, cmap='gray')
plt.title(str(epochs) + 'img' + str(i + 1))
plt.show()
torch.save(rbm_mnist, 'RBM_one.pth')
