pyrox.dev
1{ lib
2, buildPythonPackage
3, fetchFromGitHub
4, pytestCheckHook
5, pytest-xdist
6, torchvision
7, pythonOlder
8, matplotlib
9, mock
10, packaging
11, torch
12, scikit-learn
13, tqdm
14}:
15
16buildPythonPackage rec {
17 pname = "ignite";
18 version = "0.4.10";
19
20 src = fetchFromGitHub {
21 owner = "pytorch";
22 repo = pname;
23 rev = "refs/tags/v${version}";
24 sha256 = "sha256-mMiEVenDBNmeXMrDSZamUpnSm+4BQEgfK89zxIaFMio=";
25 };
26
27 checkInputs = [ pytestCheckHook matplotlib mock pytest-xdist torchvision ];
28 propagatedBuildInputs = [ packaging torch scikit-learn tqdm ];
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_gpu_info" # needs pynvml
55 "test_integration"
56 "test_output_handler" # needs mlflow
57 "test_pbar" # slight output differences
58 "test_setup_clearml_logging"
59 "test_setup_neptune"
60 "test_setup_plx"
61 "test_write_results"
62 "trains"
63 "visdom"
64 ];
65
66 meta = with lib; {
67 description = "High-level training library for PyTorch";
68 homepage = "https://pytorch.org/ignite";
69 license = licenses.bsd3;
70 maintainers = [ maintainers.bcdarwin ];
71 };
72}