Hello everyone,
I am starting to build a physical simulation framework, that will need optimisation in the future.
It is porting a matlab set of routines that do mostly linear algebra. I have already ported one part in pytorch, directly replacing numpy functionality. Beeing able to send workflows to GPU is beyond great (100x speed up for our case). But now, I am refactoring to be able to do gradient descent on the objects and classes defined on the framework.
I want some help on understanding the best practices on how to do so. I am using a lot of code from Kornia and Kaolin as my problem is descriptible in 3D meshes.
I have some questions:
- How to store class states that are goind to be diferentiated afterwards. A minimal example would be something like this:
class Foo:
def __init__(self, x: Tensor):
self.x = x
def transform(self, p: Tensor):
self.x = self.x + p
def loss(self):
return self.x.norm()**2
In this class, the sate is x
, and I would transform this state using built-in methods as transform
and would like to be able to compute a loss of his new transformed state. So I do this:
test = Foo(Tensor([1.,3.,5.]))
p = torch.tensor([-3., 1., 4. ], requires_grad=True)
And then a minimal training loop:
#train loop
learning_rate = 1e-5
for t in range(10):
test.transform(p)
loss = test.loss()
loss.backward()
print(f'loss: {loss}')
with torch.no_grad():
print(f' p:{p},\n grad {p.grad}')
p -= learning_rate * p.grad
if p.grad.norm()<1e-4:
break
p.grad.zero_()
But this is no working as expected, probably because I am overwriting the state x
with a modified x+p
in each state. In my case transformations are a little more complicated (rotations, scaling and translates).
- How would you proceed? (nn.Module extension?)
- Is there tutorials on this?
Thanks you =)