Hello there. I’m wanted to know if what I’m doing seems sensible or correct? I’ve been using Tensorflow for some time and I’m looking at PyTorch as an alternative. In a nutshell, I have an image of a known solid object (a torus) represented as a point set. This is rotated around two axes (X and Z) and rendered to an image with a gaussian blur.
*I’d post the image here but this forum only allows one image per post (which seems a little low) *
The idea now, is given this set of points, can a neural net find the angles Rx and Rz that created this image?
My plan was initially to render an image, compare this generated image to the baseline and generate a loss that way. However, this didn’t work. I realised that one could transform a point from 3D to 2D pixel coordinates easily enough but then assigning this to an image and generating a loss function won’t work as you can’t get a derivative from an index-and-assign operation. Makes sense.
Instead, I found the function grid_sample and decided that I could generate pixel coordinates and check whether or not our generated pixel position has ‘hit’ a valid area (i.e is the sample white (1.o) or not).
I wrote some basic code to see how good this loss function is:
def test(points, rot_mat, mask) :
  
  size = (128, 128) # N x C x H x W apparently
  base = test_compare()
  base = base.expand((1,1,size[0],size[1]))
  learning_rate = 0.1
  #x_rot = torch.tensor([0.0], dtype=torch.float32, requires_grad=True)
  #y_rot = torch.tensor([0.0], dtype=torch.float32, requires_grad=True)
  #z_rot = torch.tensor([0.0], dtype=torch.float32, requires_grad=True)
  
  ndc_mat = gen_ndc(size)
  #rot_mat = gen_rot(x_rot, y_rot, z_rot)
  trans_mat = gen_trans(0.0, 0.0, 2.0)
  proj_mat = gen_perspective(math.radians(90), 1.0, 1.0, 10.0)
  rot_mat.retain_grad() # Need this because it's not a leaf node?!?!

  for i in range(200):
    
    model_mat = torch.matmul(trans_mat, rot_mat)
    o = torch.matmul(model_mat, points)
    q = torch.matmul(proj_mat, o)
    # Divide through by W seems to work with a handy mask and sum
    w = q * mask
    w = torch.sum(w,1,keepdim=True)
    r = q / w
  
    s = r.narrow(1, 0, 2).reshape(1,1,-1,2)
    output = F.grid_sample(base, s)
    gauss_point = torch.tensor([[[[1.0]]]], requires_grad=True,\
      dtype=torch.float32)

    criterion = torch.nn.MSELoss()
    loss = criterion(output, gauss_point)
    loss.backward(retain_graph=True)

    with torch.no_grad():
      rot_mat -= learning_rate * rot_mat.grad
      #print("rot_mat", rot_mat)
      #loss.backward(retain_graph=True)
      #print("loss", rot_mat._grad)
      #tn = grad(loss,rot_mat)
      rot_mat.grad.zero_()
      splatted = splat(points, model_mat, proj_mat, ndc_mat, size)
      img = Image.fromarray(np.uint8(splatted.detach().numpy() * 255))
      img.save("torch" + str(i).zfill(3) + ".jpg", "JPEG")
  return rot_mat


For brevity I haven’t included the other functions (but maybe I’ll post the entire thing somewhere else).
The final result looks a little like this:

Ultimately, this hasn’t really worked. One clue is the final rotation matrix looks like this:
tensor([[ 0.4880, -0.1858, -0.0320, -1.2093],
        [ 0.1985,  1.2921, -0.0137, -0.1724],
        [ 0.0349, -0.0112,  0.8019,  0.1691],
        [ 0.0697, -0.0224, -0.3963,  1.3382]], grad_fn=<MmBackward>)


Clearly that’s not correct as the last row should be 0,0,0,1 in the ideal case. I’ve clearly messed up with the maths somewhere. I realise this is quite a tricky problem and I’ve barely scratched the surface. I suppose my first question would be:
loss.backward(retain_graph=True)

Why do I need retain_graph=True here? Without it, I get no .grad attached to my rot_mat. I thought that rot_mat would be a leaf node and would therefore have it’s gradients worked out and ready for me to apply?
Cheers
Ben