diff --git a/pkgs/development/python-modules/pytorch3d/default.nix b/pkgs/development/python-modules/pytorch3d/default.nix new file mode 100644 index 000000000000..2a92a51ec8bf --- /dev/null +++ b/pkgs/development/python-modules/pytorch3d/default.nix @@ -0,0 +1,71 @@ +{ + lib, + buildPythonPackage, + fetchFromGitHub, + setuptools, + wheel, + torch, + iopath, + cudaPackages, + config, + cudaSupport ? config.cudaSupport, +}: + +assert cudaSupport -> torch.cudaSupport; + +buildPythonPackage rec { + pname = "pytorch3d"; + version = "0.7.8"; + pyproject = true; + + src = fetchFromGitHub { + owner = "facebookresearch"; + repo = "pytorch3d"; + rev = "V${version}"; + hash = "sha256-DEEWWfjwjuXGc0WQInDTmtnWSIDUifyByxdg7hpdHlo="; + }; + + nativeBuildInputs = lib.optionals cudaSupport [ cudaPackages.cuda_nvcc ]; + build-system = [ + setuptools + wheel + ]; + dependencies = [ + torch + iopath + ]; + buildInputs = [ (lib.getOutput "cxxdev" torch) ]; + + env = + { + FORCE_CUDA = cudaSupport; + } + // lib.optionalAttrs cudaSupport { + TORCH_CUDA_ARCH_LIST = "${lib.concatStringsSep ";" torch.cudaCapabilities}"; + }; + + pythonImportsCheck = [ "pytorch3d" ]; + + passthru.tests.rotations-cuda = + cudaPackages.writeGpuTestPython { libraries = ps: [ ps.pytorch3d ]; } + '' + import pytorch3d.transforms as p3dt + + M = p3dt.random_rotations(n=10, device="cuda") + assert "cuda" in M.device.type + angles = p3dt.matrix_to_euler_angles(M, "XYZ") + assert "cuda" in angles.device.type + assert angles.shape == (10, 3), angles.shape + print(angles) + ''; + + meta = { + description = "FAIR's library of reusable components for deep learning with 3D data"; + homepage = "https://github.com/facebookresearch/pytorch3d"; + license = lib.licenses.bsd3; + maintainers = with lib.maintainers; [ + pbsds + SomeoneSerge + ]; + }; +} diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix index 92ae3b9d02da..3ab262600b43 100644 --- a/pkgs/top-level/python-packages.nix +++ b/pkgs/top-level/python-packages.nix @@ -10507,6 +10507,8 @@ self: super: with self; { pytomorrowio = callPackage ../development/python-modules/pytomorrowio { }; + pytorch3d = callPackage ../development/python-modules/pytorch3d { }; + pytouchlinesl = callPackage ../development/python-modules/pytouchlinesl { }; pyuca = callPackage ../development/python-modules/pyuca { };