pkgs/development/python-modules/ignite/default.nix

   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
  16 buildPythonPackage rec {
  17   pname = "ignite";
  18   version = "0.4.13";
  19
  20   src = fetchFromGitHub {
  21     owner = "pytorch";
  22     repo = pname;
  23     rev = "refs/tags/v${version}";
  24     hash = "sha256-+olp+zphcHMvcGKHt0JhxXza1wd7UiydwVFnSQ310Vg=";
  25   };
  26
  27   nativeCheckInputs = [ pytestCheckHook matplotlib mock pytest-xdist torchvision ];
  28   propagatedBuildInputs = [ packaging torch scikit-learn tqdm ];
  29
  30   # runs successfully 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 }