Spaces:
Sleeping
Sleeping
| # Copyright (c) 2020-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | |
| # | |
| # NVIDIA CORPORATION, its affiliates and licensors retain all intellectual | |
| # property and proprietary rights in and to this material, related | |
| # documentation and any modifications thereto. Any use, reproduction, | |
| # disclosure or distribution of this material and related documentation | |
| # without an express license agreement from NVIDIA CORPORATION or | |
| # its affiliates is strictly prohibited. | |
| import torch | |
| import os | |
| import sys | |
| sys.path.insert(0, os.path.join(sys.path[0], '../..')) | |
| import renderutils as ru | |
| RES = 4 | |
| DTYPE = torch.float32 | |
| def relative_loss(name, ref, cuda): | |
| ref = ref.float() | |
| cuda = cuda.float() | |
| print(name, torch.max(torch.abs(ref - cuda) / torch.abs(ref + 1e-7)).item()) | |
| def test_normal(): | |
| pos_cuda = torch.rand(1, RES, RES, 3, dtype=DTYPE, device='cuda', requires_grad=True) | |
| pos_ref = pos_cuda.clone().detach().requires_grad_(True) | |
| view_pos_cuda = torch.rand(1, RES, RES, 3, dtype=DTYPE, device='cuda', requires_grad=True) | |
| view_pos_ref = view_pos_cuda.clone().detach().requires_grad_(True) | |
| perturbed_nrm_cuda = torch.rand(1, RES, RES, 3, dtype=DTYPE, device='cuda', requires_grad=True) | |
| perturbed_nrm_ref = perturbed_nrm_cuda.clone().detach().requires_grad_(True) | |
| smooth_nrm_cuda = torch.rand(1, RES, RES, 3, dtype=DTYPE, device='cuda', requires_grad=True) | |
| smooth_nrm_ref = smooth_nrm_cuda.clone().detach().requires_grad_(True) | |
| smooth_tng_cuda = torch.rand(1, RES, RES, 3, dtype=DTYPE, device='cuda', requires_grad=True) | |
| smooth_tng_ref = smooth_tng_cuda.clone().detach().requires_grad_(True) | |
| geom_nrm_cuda = torch.rand(1, RES, RES, 3, dtype=DTYPE, device='cuda', requires_grad=True) | |
| geom_nrm_ref = geom_nrm_cuda.clone().detach().requires_grad_(True) | |
| target = torch.rand(1, RES, RES, 3, dtype=DTYPE, device='cuda') | |
| ref = ru.prepare_shading_normal(pos_ref, view_pos_ref, perturbed_nrm_ref, smooth_nrm_ref, smooth_tng_ref, geom_nrm_ref, True, use_python=True) | |
| ref_loss = torch.nn.MSELoss()(ref, target) | |
| ref_loss.backward() | |
| cuda = ru.prepare_shading_normal(pos_cuda, view_pos_cuda, perturbed_nrm_cuda, smooth_nrm_cuda, smooth_tng_cuda, geom_nrm_cuda, True) | |
| cuda_loss = torch.nn.MSELoss()(cuda, target) | |
| cuda_loss.backward() | |
| print("-------------------------------------------------------------") | |
| print(" bent normal") | |
| print("-------------------------------------------------------------") | |
| relative_loss("res:", ref, cuda) | |
| relative_loss("pos:", pos_ref.grad, pos_cuda.grad) | |
| relative_loss("view_pos:", view_pos_ref.grad, view_pos_cuda.grad) | |
| relative_loss("perturbed_nrm:", perturbed_nrm_ref.grad, perturbed_nrm_cuda.grad) | |
| relative_loss("smooth_nrm:", smooth_nrm_ref.grad, smooth_nrm_cuda.grad) | |
| relative_loss("smooth_tng:", smooth_tng_ref.grad, smooth_tng_cuda.grad) | |
| relative_loss("geom_nrm:", geom_nrm_ref.grad, geom_nrm_cuda.grad) | |
| def test_schlick(): | |
| f0_cuda = torch.rand(1, RES, RES, 3, dtype=DTYPE, device='cuda', requires_grad=True) | |
| f0_ref = f0_cuda.clone().detach().requires_grad_(True) | |
| f90_cuda = torch.rand(1, RES, RES, 3, dtype=DTYPE, device='cuda', requires_grad=True) | |
| f90_ref = f90_cuda.clone().detach().requires_grad_(True) | |
| cosT_cuda = torch.rand(1, RES, RES, 1, dtype=DTYPE, device='cuda', requires_grad=True) * 2.0 | |
| cosT_cuda = cosT_cuda.clone().detach().requires_grad_(True) | |
| cosT_ref = cosT_cuda.clone().detach().requires_grad_(True) | |
| target = torch.rand(1, RES, RES, 3, dtype=DTYPE, device='cuda') | |
| ref = ru._fresnel_shlick(f0_ref, f90_ref, cosT_ref, use_python=True) | |
| ref_loss = torch.nn.MSELoss()(ref, target) | |
| ref_loss.backward() | |
| cuda = ru._fresnel_shlick(f0_cuda, f90_cuda, cosT_cuda) | |
| cuda_loss = torch.nn.MSELoss()(cuda, target) | |
| cuda_loss.backward() | |
| print("-------------------------------------------------------------") | |
| print(" Fresnel shlick") | |
| print("-------------------------------------------------------------") | |
| relative_loss("res:", ref, cuda) | |
| relative_loss("f0:", f0_ref.grad, f0_cuda.grad) | |
| relative_loss("f90:", f90_ref.grad, f90_cuda.grad) | |
| relative_loss("cosT:", cosT_ref.grad, cosT_cuda.grad) | |
| def test_ndf_ggx(): | |
| alphaSqr_cuda = torch.rand(1, RES, RES, 1, dtype=DTYPE, device='cuda', requires_grad=True) | |
| alphaSqr_cuda = alphaSqr_cuda.clone().detach().requires_grad_(True) | |
| alphaSqr_ref = alphaSqr_cuda.clone().detach().requires_grad_(True) | |
| cosT_cuda = torch.rand(1, RES, RES, 1, dtype=DTYPE, device='cuda', requires_grad=True) * 3.0 - 1 | |
| cosT_cuda = cosT_cuda.clone().detach().requires_grad_(True) | |
| cosT_ref = cosT_cuda.clone().detach().requires_grad_(True) | |
| target = torch.rand(1, RES, RES, 1, dtype=DTYPE, device='cuda') | |
| ref = ru._ndf_ggx(alphaSqr_ref, cosT_ref, use_python=True) | |
| ref_loss = torch.nn.MSELoss()(ref, target) | |
| ref_loss.backward() | |
| cuda = ru._ndf_ggx(alphaSqr_cuda, cosT_cuda) | |
| cuda_loss = torch.nn.MSELoss()(cuda, target) | |
| cuda_loss.backward() | |
| print("-------------------------------------------------------------") | |
| print(" Ndf GGX") | |
| print("-------------------------------------------------------------") | |
| relative_loss("res:", ref, cuda) | |
| relative_loss("alpha:", alphaSqr_ref.grad, alphaSqr_cuda.grad) | |
| relative_loss("cosT:", cosT_ref.grad, cosT_cuda.grad) | |
| def test_lambda_ggx(): | |
| alphaSqr_cuda = torch.rand(1, RES, RES, 1, dtype=DTYPE, device='cuda', requires_grad=True) | |
| alphaSqr_ref = alphaSqr_cuda.clone().detach().requires_grad_(True) | |
| cosT_cuda = torch.rand(1, RES, RES, 1, dtype=DTYPE, device='cuda', requires_grad=True) * 3.0 - 1 | |
| cosT_cuda = cosT_cuda.clone().detach().requires_grad_(True) | |
| cosT_ref = cosT_cuda.clone().detach().requires_grad_(True) | |
| target = torch.rand(1, RES, RES, 1, dtype=DTYPE, device='cuda') | |
| ref = ru._lambda_ggx(alphaSqr_ref, cosT_ref, use_python=True) | |
| ref_loss = torch.nn.MSELoss()(ref, target) | |
| ref_loss.backward() | |
| cuda = ru._lambda_ggx(alphaSqr_cuda, cosT_cuda) | |
| cuda_loss = torch.nn.MSELoss()(cuda, target) | |
| cuda_loss.backward() | |
| print("-------------------------------------------------------------") | |
| print(" Lambda GGX") | |
| print("-------------------------------------------------------------") | |
| relative_loss("res:", ref, cuda) | |
| relative_loss("alpha:", alphaSqr_ref.grad, alphaSqr_cuda.grad) | |
| relative_loss("cosT:", cosT_ref.grad, cosT_cuda.grad) | |
| def test_masking_smith(): | |
| alphaSqr_cuda = torch.rand(1, RES, RES, 1, dtype=DTYPE, device='cuda', requires_grad=True) | |
| alphaSqr_ref = alphaSqr_cuda.clone().detach().requires_grad_(True) | |
| cosThetaI_cuda = torch.rand(1, RES, RES, 1, dtype=DTYPE, device='cuda', requires_grad=True) | |
| cosThetaI_ref = cosThetaI_cuda.clone().detach().requires_grad_(True) | |
| cosThetaO_cuda = torch.rand(1, RES, RES, 1, dtype=DTYPE, device='cuda', requires_grad=True) | |
| cosThetaO_ref = cosThetaO_cuda.clone().detach().requires_grad_(True) | |
| target = torch.rand(1, RES, RES, 1, dtype=DTYPE, device='cuda') | |
| ref = ru._masking_smith(alphaSqr_ref, cosThetaI_ref, cosThetaO_ref, use_python=True) | |
| ref_loss = torch.nn.MSELoss()(ref, target) | |
| ref_loss.backward() | |
| cuda = ru._masking_smith(alphaSqr_cuda, cosThetaI_cuda, cosThetaO_cuda) | |
| cuda_loss = torch.nn.MSELoss()(cuda, target) | |
| cuda_loss.backward() | |
| print("-------------------------------------------------------------") | |
| print(" Smith masking term") | |
| print("-------------------------------------------------------------") | |
| relative_loss("res:", ref, cuda) | |
| relative_loss("alpha:", alphaSqr_ref.grad, alphaSqr_cuda.grad) | |
| relative_loss("cosThetaI:", cosThetaI_ref.grad, cosThetaI_cuda.grad) | |
| relative_loss("cosThetaO:", cosThetaO_ref.grad, cosThetaO_cuda.grad) | |
| def test_lambert(): | |
| normals_cuda = torch.rand(1, RES, RES, 3, dtype=DTYPE, device='cuda', requires_grad=True) | |
| normals_ref = normals_cuda.clone().detach().requires_grad_(True) | |
| wi_cuda = torch.rand(1, RES, RES, 3, dtype=DTYPE, device='cuda', requires_grad=True) | |
| wi_ref = wi_cuda.clone().detach().requires_grad_(True) | |
| target = torch.rand(1, RES, RES, 1, dtype=DTYPE, device='cuda') | |
| ref = ru.lambert(normals_ref, wi_ref, use_python=True) | |
| ref_loss = torch.nn.MSELoss()(ref, target) | |
| ref_loss.backward() | |
| cuda = ru.lambert(normals_cuda, wi_cuda) | |
| cuda_loss = torch.nn.MSELoss()(cuda, target) | |
| cuda_loss.backward() | |
| print("-------------------------------------------------------------") | |
| print(" Lambert") | |
| print("-------------------------------------------------------------") | |
| relative_loss("res:", ref, cuda) | |
| relative_loss("nrm:", normals_ref.grad, normals_cuda.grad) | |
| relative_loss("wi:", wi_ref.grad, wi_cuda.grad) | |
| def test_frostbite(): | |
| normals_cuda = torch.rand(1, RES, RES, 3, dtype=DTYPE, device='cuda', requires_grad=True) | |
| normals_ref = normals_cuda.clone().detach().requires_grad_(True) | |
| wi_cuda = torch.rand(1, RES, RES, 3, dtype=DTYPE, device='cuda', requires_grad=True) | |
| wi_ref = wi_cuda.clone().detach().requires_grad_(True) | |
| wo_cuda = torch.rand(1, RES, RES, 3, dtype=DTYPE, device='cuda', requires_grad=True) | |
| wo_ref = wo_cuda.clone().detach().requires_grad_(True) | |
| rough_cuda = torch.rand(1, RES, RES, 1, dtype=DTYPE, device='cuda', requires_grad=True) | |
| rough_ref = rough_cuda.clone().detach().requires_grad_(True) | |
| target = torch.rand(1, RES, RES, 1, dtype=DTYPE, device='cuda') | |
| ref = ru.frostbite_diffuse(normals_ref, wi_ref, wo_ref, rough_ref, use_python=True) | |
| ref_loss = torch.nn.MSELoss()(ref, target) | |
| ref_loss.backward() | |
| cuda = ru.frostbite_diffuse(normals_cuda, wi_cuda, wo_cuda, rough_cuda) | |
| cuda_loss = torch.nn.MSELoss()(cuda, target) | |
| cuda_loss.backward() | |
| print("-------------------------------------------------------------") | |
| print(" Frostbite") | |
| print("-------------------------------------------------------------") | |
| relative_loss("res:", ref, cuda) | |
| relative_loss("nrm:", normals_ref.grad, normals_cuda.grad) | |
| relative_loss("wo:", wo_ref.grad, wo_cuda.grad) | |
| relative_loss("wi:", wi_ref.grad, wi_cuda.grad) | |
| relative_loss("rough:", rough_ref.grad, rough_cuda.grad) | |
| def test_pbr_specular(): | |
| col_cuda = torch.rand(1, RES, RES, 3, dtype=DTYPE, device='cuda', requires_grad=True) | |
| col_ref = col_cuda.clone().detach().requires_grad_(True) | |
| nrm_cuda = torch.rand(1, RES, RES, 3, dtype=DTYPE, device='cuda', requires_grad=True) | |
| nrm_ref = nrm_cuda.clone().detach().requires_grad_(True) | |
| wi_cuda = torch.rand(1, RES, RES, 3, dtype=DTYPE, device='cuda', requires_grad=True) | |
| wi_ref = wi_cuda.clone().detach().requires_grad_(True) | |
| wo_cuda = torch.rand(1, RES, RES, 3, dtype=DTYPE, device='cuda', requires_grad=True) | |
| wo_ref = wo_cuda.clone().detach().requires_grad_(True) | |
| alpha_cuda = torch.rand(1, RES, RES, 1, dtype=DTYPE, device='cuda', requires_grad=True) | |
| alpha_ref = alpha_cuda.clone().detach().requires_grad_(True) | |
| target = torch.rand(1, RES, RES, 3, dtype=DTYPE, device='cuda') | |
| ref = ru.pbr_specular(col_ref, nrm_ref, wo_ref, wi_ref, alpha_ref, use_python=True) | |
| ref_loss = torch.nn.MSELoss()(ref, target) | |
| ref_loss.backward() | |
| cuda = ru.pbr_specular(col_cuda, nrm_cuda, wo_cuda, wi_cuda, alpha_cuda) | |
| cuda_loss = torch.nn.MSELoss()(cuda, target) | |
| cuda_loss.backward() | |
| print("-------------------------------------------------------------") | |
| print(" Pbr specular") | |
| print("-------------------------------------------------------------") | |
| relative_loss("res:", ref, cuda) | |
| if col_ref.grad is not None: | |
| relative_loss("col:", col_ref.grad, col_cuda.grad) | |
| if nrm_ref.grad is not None: | |
| relative_loss("nrm:", nrm_ref.grad, nrm_cuda.grad) | |
| if wi_ref.grad is not None: | |
| relative_loss("wi:", wi_ref.grad, wi_cuda.grad) | |
| if wo_ref.grad is not None: | |
| relative_loss("wo:", wo_ref.grad, wo_cuda.grad) | |
| if alpha_ref.grad is not None: | |
| relative_loss("alpha:", alpha_ref.grad, alpha_cuda.grad) | |
| def test_pbr_bsdf(bsdf): | |
| kd_cuda = torch.rand(1, RES, RES, 3, dtype=DTYPE, device='cuda', requires_grad=True) | |
| kd_ref = kd_cuda.clone().detach().requires_grad_(True) | |
| arm_cuda = torch.rand(1, RES, RES, 3, dtype=DTYPE, device='cuda', requires_grad=True) | |
| arm_ref = arm_cuda.clone().detach().requires_grad_(True) | |
| pos_cuda = torch.rand(1, RES, RES, 3, dtype=DTYPE, device='cuda', requires_grad=True) | |
| pos_ref = pos_cuda.clone().detach().requires_grad_(True) | |
| nrm_cuda = torch.rand(1, RES, RES, 3, dtype=DTYPE, device='cuda', requires_grad=True) | |
| nrm_ref = nrm_cuda.clone().detach().requires_grad_(True) | |
| view_cuda = torch.rand(1, RES, RES, 3, dtype=DTYPE, device='cuda', requires_grad=True) | |
| view_ref = view_cuda.clone().detach().requires_grad_(True) | |
| light_cuda = torch.rand(1, RES, RES, 3, dtype=DTYPE, device='cuda', requires_grad=True) | |
| light_ref = light_cuda.clone().detach().requires_grad_(True) | |
| target = torch.rand(1, RES, RES, 3, dtype=DTYPE, device='cuda') | |
| ref = ru.pbr_bsdf(kd_ref, arm_ref, pos_ref, nrm_ref, view_ref, light_ref, use_python=True, bsdf=bsdf) | |
| ref_loss = torch.nn.MSELoss()(ref, target) | |
| ref_loss.backward() | |
| cuda = ru.pbr_bsdf(kd_cuda, arm_cuda, pos_cuda, nrm_cuda, view_cuda, light_cuda, bsdf=bsdf) | |
| cuda_loss = torch.nn.MSELoss()(cuda, target) | |
| cuda_loss.backward() | |
| print("-------------------------------------------------------------") | |
| print(" Pbr BSDF") | |
| print("-------------------------------------------------------------") | |
| relative_loss("res:", ref, cuda) | |
| if kd_ref.grad is not None: | |
| relative_loss("kd:", kd_ref.grad, kd_cuda.grad) | |
| if arm_ref.grad is not None: | |
| relative_loss("arm:", arm_ref.grad, arm_cuda.grad) | |
| if pos_ref.grad is not None: | |
| relative_loss("pos:", pos_ref.grad, pos_cuda.grad) | |
| if nrm_ref.grad is not None: | |
| relative_loss("nrm:", nrm_ref.grad, nrm_cuda.grad) | |
| if view_ref.grad is not None: | |
| relative_loss("view:", view_ref.grad, view_cuda.grad) | |
| if light_ref.grad is not None: | |
| relative_loss("light:", light_ref.grad, light_cuda.grad) | |
| test_normal() | |
| test_schlick() | |
| test_ndf_ggx() | |
| test_lambda_ggx() | |
| test_masking_smith() | |
| test_lambert() | |
| test_frostbite() | |
| test_pbr_specular() | |
| test_pbr_bsdf('lambert') | |
| test_pbr_bsdf('frostbite') | |