change gpu install pipeline

deps/build.jl

@@ -26,12 +26,6 @@ for pkg in ["zip", "unzip", "make", "cmake", "tensorflow=$tf_ver", "tensorflow-p
  Conda.add(pkg)
 end
 
-if Sys.islinux() && haskey(ENV, "GPU") && !("tensorflow-gpu" in pkgs)
- @info "Add tensorflow-gpu"
- # Conda.add("tensorflow-gpu=$tf_ver")
- run(`$PIP install tensorflow-gpu==$tf_ver`)
-end
-
 @info "Fix libtensorflow_framework.so..."
 if haskey(ENV, "LD_LIBRARY_PATH")
  run(setenv(`$PYTHON build.py`, "LD_LIBRARY_PATH"=>ENV["LD_LIBRARY_PATH"]*":$(Conda.LIBDIR)"))
@@ -56,49 +50,10 @@ function install_custom_op_dependency()
  run(`$UNZIP $LIBDIR/eigen.zip`)
  mv("eigen-eigen-323c052e1731", "$LIBDIR/eigen3", force=true)
  end
-
- # Install Torch library
- #=
- if Sys.isapple()
- if !isfile("$LIBDIR/libtorch.zip")
- download("https://download.pytorch.org/libtorch/cpu/libtorch-macos-latest.zip","$LIBDIR/libtorch.zip")
- end
- if !isdir("$LIBDIR/libtorch")
- run(`$UNZIP $LIBDIR/libtorch.zip`)
- mv("libtorch", "$LIBDIR/libtorch", force=true)
- if !isdir("$LIBDIR/libtorch/lib/")
- mkdir("$LIBDIR/libtorch/lib/")
- end
- download("https://github.com/intel/mkl-dnn/releases/download/v0.19/mklml_mac_2019.0.5.20190502.tgz","$LIBDIR/mklml_mac_2019.0.5.20190502.tgz")
- run(`tar -xvzf $LIBDIR/mklml_mac_2019.0.5.20190502.tgz`)
- mv("mklml_mac_2019.0.5.20190502/lib/libiomp5.dylib","$LIBDIR/libtorch/lib/libiomp5.dylib", force=true)
- mv("mklml_mac_2019.0.5.20190502/lib/libmklml.dylib","$LIBDIR/libtorch/lib/libmklml.dylib", force=true)
- rm("mklml_mac_2019.0.5.20190502/", force=true, recursive=true)
- end
- elseif Sys.islinux()
- if !isfile("$LIBDIR/libtorch.zip")
- download("https://download.pytorch.org/libtorch/cpu/libtorch-shared-with-deps-latest.zip","$LIBDIR/libtorch.zip")
- end
- if !isdir("$LIBDIR/libtorch")
- run(`$UNZIP $LIBDIR/libtorch.zip`)
- mv("libtorch", "$LIBDIR/libtorch")
- if !isdir("$LIBDIR/libtorch/lib/")
- mkdir("$LIBDIR/libtorch/lib/")
- end
- download("https://github.com/intel/mkl-dnn/releases/download/v0.19/mklml_lnx_2019.0.5.20190502.tgz","$LIBDIR/mklml_lnx_2019.0.5.20190502.tgz")
- run(`tar -xvzf $LIBDIR/mklml_lnx_2019.0.5.20190502.tgz`)
- mv("mklml_lnx_2019.0.5.20190502/lib/libiomp5.so", "$LIBDIR/libtorch/lib/libiomp5.so", force=true)
- mv("mklml_lnx_2019.0.5.20190502/lib/libmklml_gnu.so", "$LIBDIR/libtorch/lib/libmklml_gnu.so", force=true)
- mv("mklml_lnx_2019.0.5.20190502/lib/libmklml_intel.so", "$LIBDIR/libtorch/lib/libmklml_intel.so", force=true)
- rm("mklml_lnx_2019.0.5.20190502/", force=true, recursive=true)
- end
- end
- =#
 end
 
 install_custom_op_dependency()
 
-
 # useful command for debug
 # readelf -p .comment libtensorflow_framework.so 
 # strings libstdc++.so.6 | grep GLIBCXX

docs/src/customop.md

@@ -156,10 +156,12 @@ It should gives you the location of `nvcc` compiler.
 For quick installation, you can try
 ```julia
 using ADCME
-ADCME.install_gpu_dependencies()
+enable_gpu()
 ```
 
-This function basically implements the following steps
+#### Manual Installation
+
+In case 
 
 - To install CUDA toolkit (if you do not have one), you can install via conda
 ```julia

src/extra.jl

@@ -4,7 +4,8 @@ xavier_init,
 load_op_and_grad,
 load_op,
 compile_op,
-test_custom_op
+test_custom_op,
+enable_gpu
 
 """
  xavier_init(size, dtype=Float64)
@@ -251,14 +252,14 @@ function test_custom_op()
  true
 end
 
-function install_gpu_dependencies()
- Conda.add("cudatoolkit", channel="anaconda")
+function enable_gpu()
+ pkgs = Conda._installed_packages()
 
- if !("tensorflow-gpu" in Conda._installed_packages())
+ if !("tensorflow-gpu" in pkgs)
  Conda.add("tensorflow-gpu=1.14")
  end
 
- if !("cudatoolkit" in Conda._installed_packages())
+ if !("cudatoolkit" in pkgs)
  Conda.add("cudatoolkit", channel="anaconda")
  end
 
@@ -272,12 +273,33 @@ function install_gpu_dependencies()
  end
  incpath = joinpath(splitdir(strip(read(`which nvcc`, String)))[1], "../include/")
  if !isdir(joinpath(gpus, "include"))
- mv(incpath, gpus)
+ cp(incpath, joinpath(gpus, "include"))
  end
 
  pth = joinpath(Conda.ROOTENV, "pkgs/cudatoolkit-10.1.168-0/lib/")
+ # compatible 
+ files = readdir(pth)
+ for f in files
+ if f[end-2:end]==".10" && !isfile(joinpath(pth, f*".0"))
+ symlink(joinpath(pth, f), joinpath(pth, f*".0"))
+ end
+ if f[end-4:end]==".10.1" && !isfile(joinpath(pth, f[1:end-2]*".0"))
+ symlink(joinpath(pth, f), joinpath(pth, f[1:end-2]*".0"))
+ end
+ end
 
  println("Run the following command in shell
 
-export LD_LIBRARY_PATH=$pth:\$LD_LIBRARY_PATH >> ~/.bashrc")
-end
+echo 'export LD_LIBRARY_PATH=$pth:\$LD_LIBRARY_PATH' >> ~/.bashrc")
+end
+
+function use_gpu(i::Union{Nothing,Int64}=nothing)
+ if isnothing(i)
+ return 
+ end
+ i = join(collect(0:i-1),',')
+ ENV["CUDA_VISIBLE_DEVICES"] = i 
+ dl = pyimport("tensorflow.python.client.device_lib")
+ return dl.list_local_devices()
+end
+