samgalanakis / pytorch3dpointnet Goto Github PK

    verts = [torch.Tensor(p).to(device) for p in pcds]

    illum = [torch.Tensor(x).to(device) for x in lx]

    trg_point_clouds = Pointclouds(points=verts, features=illum)

    src_point_clouds = Pointclouds(points=verts)

    # We want to learn to estimate the source point cloud vertices energy
    # The shape of the estimated energy values is equal to the total number of vertices in trg_mesh
    _illum = torch.full(trg_point_clouds.features_packed().shape, 0.0, device=device, requires_grad=True)

    # The optimizer
    optimizer = torch.optim.SGD([_illum], lr=1.0, momentum=0.9)

    n_iter = 2000
    epoch = tqdm(range(n_iter))

    loss_f = torch.nn.L1Loss()
    losses = []

    for i in epoch:
        # Initialize optimizer
        optimizer.zero_grad()

        # Update of features
        src_point_clouds.features_packed = _illum

        # We compare the two sets of pointclouds by computing the L1 loss
        loss = loss_f(trg_point_clouds, src_point_clouds)

        # Print the losses
        epoch.set_description('total_loss = %.6f' % loss)

        # Save the losses for plotting
        losses.append(loss)

        # Optimization step
        loss.backward()
        optimizer.step()