MF2.py

#basic libary
import torch
import numpy as np
import matplotlib.pyplot as plt
import torch.nn as nn
import torch.optim as optim
import pandas as pd

##Define Function and class to be used

prox_plus = nn.Threshold(0,0) ## to make all output postive 

class NMF3(nn.Module):##Model for task 3
    def __init__(self, u, v, t, d):
        super(NMF3, self).__init__()
        self.U = nn.Parameter(torch.rand(u, d, requires_grad=True))
        self.V = nn.Parameter(torch.rand(d, v, requires_grad=True))
        self.T = nn.Parameter(torch.rand(t, d, requires_grad=True))

    def forward(self):
#        return (self.AA>0.5).float()
        res1 = prox_plus(torch.matmul(self.U,self.V))
        res2 = prox_plus(torch.matmul(self.T,self.V))
        return res1,res2



#Set dimension parpemeter and d
u = Xtorch.shape[0]
v = Xtorch.shape[1]
t = Ytorch.shape[0]
d = 50

task3 = NMF3(u,v,t,d)
n_epoch =100
loss_fn = nn.MSELoss(reduction='sum')

task3loss=[]
optimizer = optim.SGD(task3.parameters(), lr=0.000001)
for epoch in range(n_epoch):
    X_,Y_ = task3()
    loss = loss_fn(X_, Xtorch) +loss_fn(Y_, Ytorch)
    task3.zero_grad() # need to clear the old gradients
    loss.backward()
    optimizer.step()
    
#    print(loss)
    task3loss.append(loss)

print('Learning curve for Task 3')
plt.plot(task3loss)
plt.ylabel('loss over time')
plt.xlabel('iteration times')
plt.show()
print('Final loss on Task 3: ')
print(task3loss[-1])