cuda112 & add pytorch to ml-notebook & add cupy

ml-notebook/environment.yml

@@ -5,7 +5,9 @@ channels:
  - conda-forge
  - nodefaults
 dependencies:
- - cudatoolkit=10
- - tensorflow=2.7.0
+ - tensorflow==*=*cuda112*
+ - pytorch==*=*cuda112*
+ - torchgeo
  - jax=0.3.0
  - jaxlib=0.1.75
+ - cupy

ml-notebook/packages.txt

@@ -12,6 +12,7 @@ aiohttp==3.8.1
 aioitertools==0.10.0
 aiosignal==1.2.0
 alembic==1.8.0
+antlr-python-runtime==4.9.3
 anyio==3.6.1
 aom==3.3.0
 apache-beam==2.39.0
@@ -95,8 +96,9 @@ configobj==5.0.6
 crcmod==1.7
 croniter==0.3.36
 cryptography==37.0.2
-cudatoolkit==10.2.89
-cudnn==7.6.5.32
+cudatoolkit==11.7.0
+cudnn==8.2.1.32
+cupy==10.5.0
 curl==7.83.1
 cycler==0.11.0
 cython==0.29.30
@@ -123,6 +125,8 @@ docutils==0.15.2
 donfig==0.7.0
 earthdata==0.2.2
 eccodes==2.26.0
+efficientnet-pytorch==0.6.3
+einops==0.4.1
 entrypoints==0.4
 eofs==1.4.0
 erddapy==1.2.1
@@ -135,6 +139,7 @@ fastavro==1.5.1
 fasteners==0.17.3
 fastjmd95==0.2.1
 fastprogress==1.0.2
+fastrlock==0.8
 ffmpeg==4.4.1
 findlibs==0.0.2
 fiona==1.8.21
@@ -266,6 +271,7 @@ kerchunk==0.0.6
 kernel-headers_linux-64==2.6.32
 keyutils==1.6.1
 kiwisolver==1.4.3
+kornia==0.6.5
 krb5==1.19.3
 lame==3.100
 lcms2==2.12
@@ -299,11 +305,10 @@ libgomp==12.1.0
 libiconv==1.16
 libkml==1.3.0
 liblapack==3.9.0
-libllvm10==10.0.1
+libllvm11==11.1.0
 libnetcdf==4.8.1
 libnghttp2==1.47.0
 libnsl==2.0.0
-libopenblas==0.3.20
 libpciaccess==0.16
 libpng==1.6.37
 libpq==14.2
@@ -331,11 +336,13 @@ libxslt==1.1.33
 libzip==1.8.0
 libzlib==1.2.12
 libzopfli==1.0.3
-llvmlite==0.36.0
+llvm-openmp==14.0.4
+llvmlite==0.38.1
 locket==1.0.0
 lxml==4.8.0
 lz4==4.0.0
 lz4-c==1.9.3
+magma==2.5.4
 mailchecker==4.1.17
 mako==1.2.0
 mapclassify==2.4.3
@@ -348,6 +355,7 @@ matplotlib-inline==0.1.3
 mercantile==1.2.1
 metpy==1.3.0
 mistune==0.8.4
+mkl==2022.1.0
 morecantile==3.1.2
 mpi==1.0
 mpi4py==3.1.3
@@ -380,16 +388,18 @@ netcdf-fortran==4.5.4
 netcdf4==1.5.8
 nettle==3.6
 networkx==2.8.3
+ninja==1.11.0
 noise==1.2.2
 notebook==6.4.12
 notebook-shim==0.1.0
 nspr==4.32
 nss==3.78
-numba==0.53.1
+numba==0.55.1
 numcodecs==0.9.1
-numpy==1.22.4
+numpy==1.21.6
 oauthlib==3.2.0
 odc-geo==0.2.0
+omegaconf==2.2.2
 openh264==2.1.1
 openjpeg==2.4.0
 openssl==1.1.1o
@@ -431,6 +441,7 @@ portalocker==2.4.0
 postgresql==14.2
 pqdm==0.1.0
 prefect==0.15.13
+pretrainedmodels==0.7.4
 progressbar2==4.0.0
 proj==9.0.0
 prometheus_client==0.14.1
@@ -453,6 +464,7 @@ pyct-core==0.4.6
 pycurl==7.45.1
 pydantic==1.9.1
 pydap==3.3.0
+pydeprecate==0.3.2
 pydot==1.4.2
 pyerfa==2.0.0.1
 pygeos==0.12.0
@@ -494,6 +506,8 @@ python-slugify==6.1.2
 python-utils==3.3.3
 python_abi==3.9
 pytools==2022.1.9
+pytorch==1.11.0
+pytorch-lightning==1.6.4
 pytz==2022.1
 pytzdata==2020.1
 pyu2f==0.1.5
@@ -523,11 +537,13 @@ scikit-learn==1.1.1
 scipy==1.8.1
 seaborn==0.11.2
 seaborn-base==0.11.2
+segmentation-models-pytorch==0.2.1
 send2trash==1.8.0
-setuptools==62.3.4
+setuptools==59.5.0
 shapely==1.8.2
 simpervisor==0.4
 six==1.16.0
+sleef==3.5.1
 smmap==3.0.5
 snappy==1.1.9
 sniffio==1.2.0
@@ -544,6 +560,7 @@ statsmodels==0.13.2
 svt-av1==0.9.1
 sysroot_linux-64==2.12
 tabulate==0.8.9
+tbb==2021.5.0
 tblib==1.7.0
 tensorboard==2.6.0
 tensorboard-data-server==0.6.0
@@ -559,11 +576,15 @@ tifffile==2022.5.4
 tiledb==2.8.3
 tiledb-py==0.14.5
 timezonefinder==5.2.0
+timm==0.4.12
 tinycss2==1.1.1
 tk==8.6.12
 toml==0.10.2
 tomli==2.0.1
 toolz==0.11.2
+torchgeo==0.2.1
+torchmetrics==0.8.2
+torchvision==0.12.0
 tornado==6.1
 tqdm==4.64.0
 traitlets==5.2.2.post1

pytorch-notebook/conda-linux-64.lock

@@ -1,6 +1,6 @@
 # Generated by conda-lock.
 # platform: linux-64
-# input_hash: 0829848aaf19cd7962949c1ebe0b99b894ffb49fc331f8a3a507a6382074c4be
+# input_hash: ed2263bbbda4634dff678be1377ea510f92ba88ce9e171c462474c26c3bf897c
 @EXPLICIT
 https://conda.anaconda.org/conda-forge/linux-64/_libgcc_mutex-0.1-conda_forge.tar.bz2#d7c89558ba9fa0495403155b64376d81
 https://conda.anaconda.org/conda-forge/linux-64/ca-certificates-2022.5.18.1-ha878542_0.tar.bz2#352e93bbe1d604002b11bbcf425bf866
@@ -59,7 +59,6 @@ https://conda.anaconda.org/conda-forge/linux-64/libffi-3.4.2-h7f98852_5.tar.bz2#
 https://conda.anaconda.org/conda-forge/linux-64/libglu-9.0.0-he1b5a44_1001.tar.bz2#8208602aec4826053c116552369a394c
 https://conda.anaconda.org/conda-forge/linux-64/libiconv-1.16-h516909a_0.tar.bz2#5c0f338a513a2943c659ae619fca9211
 https://conda.anaconda.org/conda-forge/linux-64/libnsl-2.0.0-h7f98852_0.tar.bz2#39b1328babf85c7c3a61636d9cd50206
-https://conda.anaconda.org/conda-forge/linux-64/libopenblas-0.3.20-pthreads_h78a6416_0.tar.bz2#9b6d0781953c9e353faee494336cc229
 https://conda.anaconda.org/conda-forge/linux-64/libpciaccess-0.16-h516909a_0.tar.bz2#b38752a755b9c64f07651fcc0084f878
 https://conda.anaconda.org/conda-forge/linux-64/libsanitizer-12.1.0-ha89aaad_16.tar.bz2#72d63459c86185f8f636772f28d6eb35
 https://conda.anaconda.org/conda-forge/linux-64/libsodium-1.0.18-h36c2ea0_1.tar.bz2#c3788462a6fbddafdb413a9f9053e58d
@@ -480,7 +479,6 @@ https://conda.anaconda.org/conda-forge/noarch/pyopenssl-22.0.0-pyhd8ed1ab_0.tar.
 https://conda.anaconda.org/conda-forge/linux-64/pyresample-1.23.0-py39h1832856_1.tar.bz2#1cf53547f0287814a0df8ac4179c0914
 https://conda.anaconda.org/conda-forge/noarch/python-box-6.0.2-pyhd8ed1ab_0.tar.bz2#b96adf535c6c58706897fb1c21c4cb54
 https://conda.anaconda.org/conda-forge/linux-64/python-geotiepoints-1.4.1-py39hd257fcd_0.tar.bz2#da51bfee40828cb40a8a13c689cbb1c0
-https://conda.anaconda.org/conda-forge/linux-64/pytorch-gpu-1.11.0-cuda112py39h68407e5_202.tar.bz2#a928246d03fbd68f821b68080329b933
 https://conda.anaconda.org/conda-forge/linux-64/scikit-learn-1.1.1-py39h4037b75_0.tar.bz2#933f279f76b1e65845c65d9966bf1533
 https://conda.anaconda.org/conda-forge/noarch/sparse-0.13.0-pyhd8ed1ab_0.tar.bz2#e9727043a1d08f997884851c9f428638
 https://conda.anaconda.org/conda-forge/noarch/starlette-0.19.1-pyhd8ed1ab_0.tar.bz2#f6618f4eab5aecbcade4a43e3c98700a

pytorch-notebook/environment.yml

@@ -5,9 +5,5 @@ channels:
  - conda-forge
  - nodefaults
 dependencies:
- - cudatoolkit=11
- # needed for pytorch to work
- # https://github.com/NVIDIA/MinkowskiEngine/issues/354#issuecomment-880555195
- - anaconda::libopenblas
- - pytorch-gpu
+ - pytorch==*=*cuda112*
  - torchgeo

pytorch-notebook/packages.txt

@@ -301,7 +301,6 @@ libllvm11==11.1.0
 libnetcdf==4.8.1
 libnghttp2==1.47.0
 libnsl==2.0.0
-libopenblas==0.3.20
 libpciaccess==0.16
 libpng==1.6.37
 libpq==14.2
@@ -498,7 +497,6 @@ python-utils==3.3.3
 python_abi==3.9
 pytools==2022.1.9
 pytorch==1.11.0
-pytorch-gpu==1.11.0
 pytorch-lightning==1.6.4
 pytz==2022.1
 pytzdata==2020.1

tests/test_ml-notebook.py

@@ -5,8 +5,13 @@
 
 packages = [
     # machine learning stuff
-    'tensorflow', 'skimage', 'sklearn',
-    'jax',
+    "tensorflow", "skimage", "sklearn",
+    "jax",
+    "pytorch_lightning",
+    "skimage",
+    "sklearn",
+    "torch",
+    "torchgeo",
     # cupy import fails unless on GPU-enabled node:
     #'cupy', #libcuda.so.1: cannot open shared object file: No such file or directory
     # Error w/conda solve of pytorch+tensorflow, so not currently installed
@@ -21,3 +26,10 @@ def test_start():
     print(os.environ)
     if os.environ.get('PANGEO_ENV') is not None:
         assert os.environ['PANGEO_ENV'] == 'ml-notebook'
+
+def test_torch_uses_mkl():
+
+    import torch
+
+    blas_info_index = torch.__config__.show().find("BLAS_INFO")
+    assert torch.__config__.show()[blas_info_index + 10 : blas_info_index + 13] == "mkl"