kuelumbus / docker-deepmd-kit

Docker image for DeePMD-kit (CPU & GPU)

Docker image for DeePMD-kit (CPU & GPU)

DeePMD-kit is a deep learning package for many-body potential energy representation, optimization, and molecular dynamics.

This docker project is set up to simplify the installation process of DeePMD-kit.

Features

• Tensorflow works with CPU & GPU
• MPI version of Lammps

Build the image

git clone https://github.com/kuelumbus/docker-deepmd-kit.git
cd deepmd-kit_docker && docker build -f Dockerfile -t deepmd-gpu .

The nvidia runtime for docker (--runtime=nvidia switch) must be installed to run the container.

Run the container

A bash executable to run dp_train | dp_test | dp_frz should look like (e.g. run as dp_train 1 in.json)

#!/bin/bash
docker run -it --rm  \
    --runtime=nvidia \
    --mount type=bind,source="$(pwd)",target=/app \
    -e CUDA_VISIBLE_DEVICES="$1" \
    deepmd-gpu:latest \
    dp_train ${@:2} 

Similar for lammps run as (lmp_mpi n_cores -in in.lammps). Make sure predictions during the lammps run are done on the CPU and not GPU by setting CUDA_VISIBLE_DEVICES="" - I am not sure if this is necessary.

#!/bin/bash
docker run -it --rm  \
    --runtime=nvidia \
    --mount type=bind,source="$(pwd)",target=/app \
    -e CUDA_VISIBLE_DEVICES="" \
    deepmd-gpu:latest \
    mpiexec --allow-run-as-root -n $1 lmp_mpi ${@:2}

The mount argument binds the current directory to /app in the docker container and thus must be executed in the directory with the input json file.