pyrox.dev
1{ lib
2, buildPythonPackage
3, fetchFromGitHub
4, pytestCheckHook
5, pytest-xdist
6, torchvision
7, pythonOlder
8, matplotlib
9, mock
10, pytorch
11, pynvml
12, scikit-learn
13, tqdm
14}:
15
16buildPythonPackage rec {
17 pname = "ignite";
18 version = "0.4.6";
19
20 src = fetchFromGitHub {
21 owner = "pytorch";
22 repo = pname;
23 rev = "v${version}";
24 sha256 = "sha256-dlKGXjUUnyYmPDilo0LQg9OkSkBnMYNgzlFLIfI0T6I=";
25 };
26
27 checkInputs = [ pytestCheckHook matplotlib mock pytest-xdist torchvision ];
28 propagatedBuildInputs = [ pytorch scikit-learn tqdm pynvml ];
29
30 # runs succesfully in 3.9, however, async isn't correctly closed so it will fail after test suite.
31 doCheck = pythonOlder "3.9";
32
33 # Some packages are not in NixPkgs; other tests try to build distributed
34 # models, which doesn't work in the sandbox.
35 # avoid tests which need special packages
36 pytestFlagsArray = [
37 "--ignore=tests/ignite/contrib/handlers/test_clearml_logger.py"
38 "--ignore=tests/ignite/contrib/handlers/test_lr_finder.py"
39 "--ignore=tests/ignite/contrib/handlers/test_trains_logger.py"
40 "--ignore=tests/ignite/metrics/nlp/test_bleu.py"
41 "--ignore=tests/ignite/metrics/nlp/test_rouge.py"
42 "--ignore=tests/ignite/metrics/gan" # requires pytorch_fid; tries to download model to $HOME
43 "--ignore=tests/ignite/metrics/test_dill.py"
44 "--ignore=tests/ignite/metrics/test_psnr.py"
45 "--ignore=tests/ignite/metrics/test_ssim.py"
46 "tests/"
47 ];
48
49 # disable tests which need specific packages
50 disabledTests = [
51 "idist"
52 "mlflow"
53 "tensorboard"
54 "test_integration"
55 "test_output_handler" # needs mlflow
56 "test_pbar" # slight output differences
57 "test_setup_clearml_logging"
58 "test_setup_neptune"
59 "test_setup_plx"
60 "test_write_results"
61 "trains"
62 "visdom"
63 ];
64
65 meta = with lib; {
66 description = "High-level training library for PyTorch";
67 homepage = "https://pytorch.org/ignite";
68 license = licenses.bsd3;
69 maintainers = [ maintainers.bcdarwin ];
70 };
71}