How to implement matrix factorization in pytorch?

1

I am trying to implement a matrix factorization algorithm in pytorch. Specifically, I have a class with matrices A,B,C and want to train A and B such that AB = C. My first try was to write the training as

for i in range(self.max_iter):
            self.optimizer.zero_grad()
            loss = ((torch.mm(self.A, self.B) - self.C)**2).mean()
            loss.backward()
            self.optimizer.step()

The variables are kept as

class nmf_solver:

    def __init__(self, A,B,C, step_size=STEP_SIZE):

        """
        solves | AB - C |^2
        """

        self.A = Variable(A.cuda(), requires_grad=True)
        self.B = Variable(B.cuda(), requires_grad=True)
        self.C = C.cuda()

        self.step_size = step_size

        self.optimizer = optim.SGD([self.A,self.B], lr=self.step_size, momentum=0.9)

But I’m met with a “RuntimeError: Trying to backward through the graph a second time, but the buffers have already been freed.”. Since there isn’t any input to the matrix factorization algorithm, I don’t really know if nn.module is appropriate to subclass.

Is there any way to make this work in pytorch?

2

The key here is how you keep A B and C. Apparently they aren’t leaf variables in your graph, but your snippet doesn’t show what they are.
It is completely legitimate to have modules without input. Their main function is to hold the trainable Parameters, in your case A and B, with the calculation.

Best regards

Thomas

3

Dear Thomas,

Thanks for your reply. I’ve edited the post to show how I keep the variables. How should I add A and B, assuming that I did subclass nn.module? Thanks!!

best,
Johan

4

Variable isn’t a thing anymore.
The mistake in your code is that you do the .cuda() after requiring the gradient. .cuda is calculation for autograd purposes.
I’d probably subclass nn.Module and self.A = nn.Parameter(A) (requires grad is automatic for nn.Parameter).
You don’t want to include your optimizer in the module itself, if only so you can do the following in the right order: a) instantiate your class b) .cuda() your instance c) set up the optimizer.

Best regards

Thomas

5

Dear Thomas,

Oh I see! Thanks so much!! :slight_smile:

best,
Johan