Simple MLE of Exponential Variables with Pytorch

mle.py

'''
  Simple MLE Estimator to demonstrate Pytorch optimization capabilities
  This program estimates the parameter of an exponential random variable
  using an MLE Estimator.
  
  Author: Marios Papachristou
'''

import torch
import numpy as np
from torch.autograd import Variable

# Draw random samples
sample = np.random.randint(10, size=(10, 1))

# Convert to Pytorch tensor
x = Variable(torch.from_numpy(sample)).type(torch.FloatTensor)

# Initialize Parameters
p = Variable(torch.rand(1), requires_grad=True)

# Initialize optimizer
optimizer = torch.optim.SGD([p], lr=0.0002)

for t in range(1500):
    optimizer.zero_grad()

    # Negative log likelihood of exponential NLL = -log (L | X_1, ..., X_n)
    NLL = - (x.size()[0] * torch.log(p) - p * torch.sum(x))

    NLL.backward()
    
    print('NLL', NLL, 'p', p)

    optimizer.step()

print('Theoretical estimation', 1 / np.mean(sample))