This repository provides an efficient PyTorch-based library for training deep models.

Related tags

Deep Learning Hammer
Overview

An Efficient Library for Training Deep Models

This repository provides an efficient PyTorch-based library for training deep models.

Installation

Make sure your Python >= 3.7, CUDA version >= 11.1, and CUDNN version >= 7.6.5.

  1. Install package requirements via conda:

    conda create -n <ENV_NAME> python=3.7  # create virtual environment with Python 3.7
    conda activate <ENV_NAME>
    pip install -r requirements/minimal.txt -f https://download.pytorch.org/whl/cu111/torch_stable.html
  2. To use video visualizer (optional), please also install ffmpeg.

    • Ubuntu: sudo apt-get install ffmpeg.
    • MacOS: brew install ffmpeg.
  3. To reduce memory footprint (optional), you can switch to either jemalloc (recommended) or tcmalloc rather than your default memory allocator.

    • jemalloc (recommended):
      • Ubuntu: sudo apt-get install libjemalloc
    • tcmalloc:
      • Ubuntu: sudo apt-get install google-perftools
  4. (optional) To speed up data loading on NVIDIA GPUs, you can install DALI, together with dill to pickle python objects. It is optional to also install CuPy for some customized operations if needed:

    pip install --extra-index-url https://developer.download.nvidia.com/compute/redist --upgrade nvidia-dali-<CUDA_VERSION>
    pip install dill
    pip install cupy  # optional, installation can be slow

    For example, on CUDA 11.1, DALI can be installed via:

    pip install --extra-index-url https://developer.download.nvidia.com/compute/redist --upgrade nvidia-dali-cuda110  # CUDA 11.1 compatible
    pip install dill
    pip install cupy  # optional, installation can be slow

Quick Demo

Train StyleGAN2 on FFHQ in Resolution of 256x256

In your Terminal, run:

./scripts/training_demos/stylegan2_ffhq256.sh <NUM_GPUS> <PATH_TO_DATA> [OPTIONS]

where

  • refers to the number of GPUs. Setting as 1 helps launch a training job on single-GPU platforms.

  • refers to the path of FFHQ dataset (in resolution of 256x256) with zip format. If running on local machines, a soft link of the data will be created under the data folder of the working directory to save disk space.

  • [OPTIONS] refers to any additional option to pass. Detailed instructions on available options can be shown via ./scripts/training_demos/stylegan2_ffhq256.sh --help .

This demo script uses stylegan2_ffhq256 as the default value of job_name, which is particularly used to identify experiments. Concretely, a directory with name job_name will be created under the root working directory (with is set as work_dirs/ by default). To prevent overwriting previous experiments, an exception will be raised to interrupt the training if the job_name directory has already existed. To change the job name, please use --job_name= option.

More Demos

Please find more training demos under ./scripts/training_demos/.

Inspect Training Results

Besides using TensorBoard to track the training process, the raw results (e.g., training losses and running time) are saved in JSON format. They can be easily inspected with the following script

import json

file_name = '
   
    /log.json'
   

data_entries = []
with open(file_name, 'r') as f:
    for line in f:
        data_entry = json.loads(line)
        data_entries.append(data_entry)

# An example of data entry
# {"Loss/D Fake": 0.4833524551040682, "Loss/D Real": 0.4966000154727226, "Loss/G": 1.1439273656869773, "Learning Rate/Discriminator": 0.002352941082790494, "Learning Rate/Generator": 0.0020000000949949026, "data time": 0.0036810599267482758, "iter time": 0.24490128830075264, "run time": 66108.140625}

Convert Pre-trained Models

See Model Conversion for details.

Prepare Datasets

See Dataset Preparation for details.

Develop

See Contributing Guide for details.

License

The project is under MIT License.

Acknowledgement

This repository originates from GenForce, with all modules carefully optimized to make it more flexible and robust for distributed training. On top of GenForce where only StyleGAN training is provided, this repository also supports training StyleGAN2 and StyleGAN3, both of which are fully reproduced. Any new method is welcome to merge into this repository! Please refer to the Develop section.

Contributors

The main contributors are listed as follows.

Member Contribution
Yujun Shen Refactor and optimize the entire codebase and reproduce start-of-the-art approaches.
Zhiyi Zhang Contribute to a number of sub-modules and functions, especially dataset related.
Dingdong Yang Contribute to DALI data loading acceleration.
Yinghao Xu Originally contribute to runner and loss functions in GenForce.
Ceyuan Yang Originally contribute to data loader in GenForce.
Jiapeng Zhu Originally contribute to evaluation metrics in GenForce.

BibTex

We open source this library to the community to facilitate the research. If you do like our work and use the codebase for your projects, please cite our work as follows.

@misc{hammer2022,
  title =        {Hammer: An Efficient Toolkit for Training Deep Models.},
  author =       {Shen, Yujun and Zhang, Zhiyi and Yang, Dingdong and Xu, Yinghao and Yang, Ceyuan and Zhu, Jiapeng},
  howpublished = {\url{https://github.com/bytedance/Hammer}},
  year =         {2022}
}
You might also like...
Efficient training of deep recommenders on cloud.
Efficient training of deep recommenders on cloud.

HybridBackend Introduction HybridBackend is a training framework for deep recommenders which bridges the gap between evolving cloud infrastructure and

ESGD-M - A stochastic non-convex second order optimizer, suitable for training deep learning models, for PyTorch

ESGD-M - A stochastic non-convex second order optimizer, suitable for training deep learning models, for PyTorch

XtremeDistil framework for distilling/compressing massive multilingual neural network models to tiny and efficient models for AI at scale

XtremeDistilTransformers for Distilling Massive Multilingual Neural Networks ACL 2020 Microsoft Research [Paper] [Video] Releasing [XtremeDistilTransf

Transformers4Rec is a flexible and efficient library for sequential and session-based recommendation, available for both PyTorch and Tensorflow.
Transformers4Rec is a flexible and efficient library for sequential and session-based recommendation, available for both PyTorch and Tensorflow.

Transformers4Rec is a flexible and efficient library for sequential and session-based recommendation, available for both PyTorch and Tensorflow.

Pytorch library for end-to-end transformer models training and serving

Pytorch library for end-to-end transformer models training and serving

TorchMultimodal is a PyTorch library for training state-of-the-art multimodal multi-task models at scale.

TorchMultimodal (Alpha Release) Introduction TorchMultimodal is a PyTorch library for training state-of-the-art multimodal multi-task models at scale.

Learning recognition/segmentation models without end-to-end training. 40%-60% less GPU memory footprint. Same training time. Better performance.
Learning recognition/segmentation models without end-to-end training. 40%-60% less GPU memory footprint. Same training time. Better performance.

InfoPro-Pytorch The Information Propagation algorithm for training deep networks with local supervision. (ICLR 2021) Revisiting Locally Supervised Lea

A library for preparing, training, and evaluating scalable deep learning hybrid recommender systems using PyTorch.
A library for preparing, training, and evaluating scalable deep learning hybrid recommender systems using PyTorch.

collie_recs Collie is a library for preparing, training, and evaluating implicit deep learning hybrid recommender systems, named after the Border Coll

A library for preparing, training, and evaluating scalable deep learning hybrid recommender systems using PyTorch.
A library for preparing, training, and evaluating scalable deep learning hybrid recommender systems using PyTorch.

collie Collie is a library for preparing, training, and evaluating implicit deep learning hybrid recommender systems, named after the Border Collie do

Comments
  • CVE-2007-4559 Patch

    CVE-2007-4559 Patch

    Patching CVE-2007-4559

    Hi, we are security researchers from the Advanced Research Center at Trellix. We have began a campaign to patch a widespread bug named CVE-2007-4559. CVE-2007-4559 is a 15 year old bug in the Python tarfile package. By using extract() or extractall() on a tarfile object without sanitizing input, a maliciously crafted .tar file could perform a directory path traversal attack. We found at least one unsantized extractall() in your codebase and are providing a patch for you via pull request. The patch essentially checks to see if all tarfile members will be extracted safely and throws an exception otherwise. We encourage you to use this patch or your own solution to secure against CVE-2007-4559. Further technical information about the vulnerability can be found in this blog.

    If you have further questions you may contact us through this projects lead researcher Kasimir Schulz.

    opened by TrellixVulnTeam 0
  • convert_model.py stylegan2

    convert_model.py stylegan2

    Thanks your great work!

    When i try to transform official weight .pkl to .pth, I found this error.

    Dealing with  /home/resefa/Hammer-main/converters/stylegan2-ffhq-config-f.pkl ...
    
    Traceback (most recent call last):
      File "convert_model.py", line 58, in <module>
        main()
      File "convert_model.py", line 54, in main
        learning_rate=args.learning_rate)
      File "/home/aiyang/resefa/Hammer-main/converters/base_converter.py", line 214, in run
        self.load_source(src_path)
      File "/home/aiyang/resefa/Hammer-main/converters/stylegan2_converter.py", line 61, in load_source
        G, D, Gs = pickle.load(f)
    ModuleNotFoundError: No module named 'torch_utils.persistence'
    

    I try to reinstall torch_utils and torch, but it didn't work. Need help!

    opened by Ruye-aa 0
Owner
Bytedance Inc.
Bytedance Inc.
DeepSpeed is a deep learning optimization library that makes distributed training easy, efficient, and effective.

DeepSpeed is a deep learning optimization library that makes distributed training easy, efficient, and effective.

Microsoft 8.4k Jan 1, 2023
Easy Parallel Library (EPL) is a general and efficient deep learning framework for distributed model training.

English | 简体中文 Easy Parallel Library Overview Easy Parallel Library (EPL) is a general and efficient library for distributed model training. Usability

Alibaba 185 Dec 21, 2022
A PyTorch implementation of "Cluster-GCN: An Efficient Algorithm for Training Deep and Large Graph Convolutional Networks" (KDD 2019).

ClusterGCN ⠀⠀ A PyTorch implementation of "Cluster-GCN: An Efficient Algorithm for Training Deep and Large Graph Convolutional Networks" (KDD 2019). A

Benedek Rozemberczki 697 Dec 27, 2022
TorchGeo is a PyTorch domain library, similar to torchvision, that provides datasets, transforms, samplers, and pre-trained models specific to geospatial data.

TorchGeo is a PyTorch domain library, similar to torchvision, that provides datasets, transforms, samplers, and pre-trained models specific to geospatial data.

Microsoft 1.3k Dec 30, 2022
Ultra-Data-Efficient GAN Training: Drawing A Lottery Ticket First, Then Training It Toughly

Ultra-Data-Efficient GAN Training: Drawing A Lottery Ticket First, Then Training It Toughly Code for this paper Ultra-Data-Efficient GAN Tra

VITA 77 Oct 5, 2022
Efficient-GlobalPointer - Pytorch Efficient GlobalPointer

引言 感谢苏神带来的模型,原文地址:https://spaces.ac.cn/archives/8877 如何运行 对应模型EfficientGlobalPoi

powerycy 40 Dec 14, 2022
TorchX: A PyTorch Extension Library for More Efficient Deep Learning

TorchX TorchX: A PyTorch Extension Library for More Efficient Deep Learning. @misc{torchx, author = {Ansheng You and Changxu Wang}, title = {T

Donny You 8 May 28, 2022
NVIDIA Merlin is an open source library providing end-to-end GPU-accelerated recommender systems, from feature engineering and preprocessing to training deep learning models and running inference in production.

NVIDIA Merlin NVIDIA Merlin is an open source library designed to accelerate recommender systems on NVIDIA’s GPUs. It enables data scientists, machine

null 419 Jan 3, 2023
[ICLR 2021] "CPT: Efficient Deep Neural Network Training via Cyclic Precision" by Yonggan Fu, Han Guo, Meng Li, Xin Yang, Yining Ding, Vikas Chandra, Yingyan Lin

CPT: Efficient Deep Neural Network Training via Cyclic Precision Yonggan Fu, Han Guo, Meng Li, Xin Yang, Yining Ding, Vikas Chandra, Yingyan Lin Accep

null 26 Oct 25, 2022