KungFu

KungFu distributed machine learning framework

Install

Make sure you have tensorflow or tensorflow-gpu python library installed.

# install
export CMAKE_BUILD_PARALLEL_LEVEL=$(nproc) # 4 seconds faster
pip3 install --no-index -U .

# FIXME: For Mac users, the following is required after the install:
# export DYLD_LIBRARY_PATH=$(python3 -c "import os; import kungfu; print(os.path.dirname(kungfu.__file__))")

To enable NCCL support

# uncomment to use your own NCCL
# export NCCL_HOME=$HOME/local/nccl

env \
    KUNGFU_USE_NCCL=1 \
    pip3 install --no-index --user -U .

Build kungfu-prun

./configure --build-tools
make

Format code

./scripts/clean-code.sh --fmt-py

Example

Download MNIST dataset (script) and run the following training script.

python3 examples/mnist_mlp.py

Build for release

# build a .whl package for release
pip3 wheel -vvv --no-index .