Optimizing mesh rotation

[xdobetter]

Hi,I am now trying to test optimization process w.r.t mesh roatate.My target is implement the same goals of psdr-sg20 tree demos by using psdr-cuda,so I refer to the code you gave in #4 and others code in #11 ,I set A=torch.tensor( [0.6] ) (and similar to psdr-sg 20), but it does't work and show "RuntimeError: Attempted to cast a Torch CPU tensor to a Enoki GPU array!"

And I have a questions for your give code in #4. 1.What does variable A mean? Is variable A the same as psdr-sg20?

and the python code is like below. I wish to get an advice,thanks.

import psdr_cuda
import torch
import time
import numpy as np
import enoki as ek
from enoki.cuda_autodiff import Float32 as FloatD ,Vector3f as Vector3fD,Matrix4f as Matrix4fD
from enoki.cuda import Float32 as FloatC,Vector3f as Vector3fC
import os
os.environ["OPENCV_IO_ENABLE_OPENEXR"]="1"
import cv2

dir_name="./res/" #my output dir

def create_output_directory(dir_out): #
    if not os.path.exists(dir_out):
        os.mkdir( dir_out )
    if not os.path.exists(dir_out + 'grad_img'):
        os.mkdir( dir_out + 'grad_img' )
    if not os.path.exists(dir_out + 'iterations'):
        os.mkdir( dir_out + 'iterations' )
    if not os.path.exists(dir_out + 'plot'):
        os.mkdir( dir_out + 'plot' )

def save_exr(img,file_name): #after reshape,use it
    out=cv2.cvtColor(img,cv2.COLOR_RGB2BGR) # OpenCV uses BGR order
    cv2.imwrite(file_name,out)

def psdr_optimize(scene_file):
    print("psdr_optimize!")
    sc = psdr_cuda.Scene()
    sc.load_file(scene_file, False)
    sc.param_map["Mesh[1]"].enable_edges = False
    ro = sc.opts
    ro.spp = 16
    ro.sppe = 16
    ro.sppse = 16
    sc.opts.log_level = 0
    num_sensors = sc.num_sensors
    sc.configure()
    axis = Vector3fD(0.0, 0.0, 1.0)
    valid_camera = 0
    print('Loading target images')
    tars = []
    #Initialize an integrator
    intergrator=psdr_cuda.DirectIntegrator()
    #render init image
    image_init=intergrator.renderD(sc,sensor_id=0)
    np_image_init=image_init.numpy().reshape(sc.opts.height,sc.opts.width,3)
    #save init image
    save_exr(np_image_init,dir_name+"init.exr")

    #target = cv2.imread(target_path+"exr/sensor_%d.exr" % valid_camera, cv2.IMREAD_UNCHANGED)
    #target = torch.from_numpy(cv2.cvtColor(target, cv2.COLOR_RGB2BGR)).float()
    #target = target.reshape((-1, 3))

    # construct a rotation matrix
    rotate_value=1.0
    P = FloatD(rotate_value) #rotate_value
    #A = [P.torch()]
    #A[0].requires_grad=True
    # A= [sc.param_map["Mesh[1]"]]
    A=torch.tensor( [0.6] )
    params = [
        {'params': A, 'lr':  0.03},
    ]

    # render target image
    mat=Matrix4fD.rotate(axis,P) #rotate matrix
    sc.param_map["Mesh[1]"].set_transform(mat) #give a rotate matrix
    image_target=intergrator.renderD(sc,sensor_id=0)
    np_image_target=image_target.numpy().reshape((sc.opts.height,sc.opts.width,3))
    # save target image
    save_exr(np_image_target,dir_name+"target.exr")

    tars.append(image_target)

    class PSDRRender(torch.autograd.Function):
        @staticmethod 
        def forward(ctx, A, iter):
            _dir_vector = FloatD(A)
            ek.set_requires_gradient(_dir_vector, A.requires_grad)
            ctx.in1 = _dir_vector
            sc.param_map['Mesh[1]'].set_transform(Matrix4fD.rotate(axis, _dir_vector))
            sc.configure()

            integrator = psdr_cuda.DirectIntegrator(1, 1, 0)
            img = integrator.renderD(sc, valid_camera)
            tar = Vector3fD(tars[0].cuda())

            loss = ek.abs(img - tar)
            loss = ek.hsum(ek.hsum(loss)) / (ro.height * ro.width * 3)
            ctx.out = loss 
            out_torch = ctx.out.torch() 
            return out_torch

        @staticmethod
        def backward(ctx, grad_out):
            ek.set_gradient(ctx.out, FloatC(grad_out))
            FloatD.backward() 
            gradA = ek.gradient(ctx.in1).torch()
            result = gradA
            del ctx.out, ctx.in1
            return (result, None)

    optimizer = torch.optim.Adam(params)
    render = PSDRRender.apply
    npass = 100 #optimize iter
    print('Starting optimization')
    loss_graph = np.zeros(npass)
    para_result = np.array([rotate_value])
    para_graph = np.zeros(npass)

    for i in range(npass):
        start = time.time()

        average_loss = 0
        average_grad = 0
        optimizer.zero_grad()
        gpass = 12
        for g in range(gpass):
            loss_temp = render(A, g)
            loss_temp.backward()
            average_loss += loss_temp.item() / float(gpass)
            average_grad += A.grad / float(gpass)
        A.grad = average_grad
        optimizer.step()
        end = time.time()

        loss_graph[i] = average_loss
        curr_para = A.cpu().detach().numpy()[0]
        para_graph[i] = abs(curr_para)
        print('i:', i, ' loss:', average_loss, ' para_diff:', curr_para, ' Time:', end - start) 
        torch.cuda.empty_cache() 
        ek.cuda_malloc_trim()

    #image_loss(npass, loss_graph, output_dir)
    #para_loss(npass, para_graph, output_dir)
if __name__=="__main__":
    create_output_directory(dir_name)
    scene_file="./data/scenes/tree.xml"
    psdr_optimize(scene_file)

[andyyankai]

You simply need something like cuda2 = torch.device('cuda:0') x = torch.tensor([....], device=cuda0)