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