1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
|
{
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
];
};
}
|
