{"id":21300102,"url":"https://github.com/roclark/torchsr","last_synced_at":"2025-07-06T23:06:06.946Z","repository":{"id":37485560,"uuid":"362215431","full_name":"roclark/torchsr","owner":"roclark","description":"An easy-to-use SRGAN-based Super Resolution tool using PyTorch","archived":false,"fork":false,"pushed_at":"2022-06-22T20:51:47.000Z","size":7612,"stargazers_count":5,"open_issues_count":1,"forks_count":1,"subscribers_count":2,"default_branch":"main","last_synced_at":"2025-06-29T18:08:00.257Z","etag":null,"topics":[],"latest_commit_sha":null,"homepage":null,"language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"apache-2.0","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/roclark.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null}},"created_at":"2021-04-27T18:38:36.000Z","updated_at":"2024-09-05T08:30:03.000Z","dependencies_parsed_at":"2022-09-15T04:41:29.734Z","dependency_job_id":null,"html_url":"https://github.com/roclark/torchsr","commit_stats":null,"previous_names":[],"tags_count":1,"template":false,"template_full_name":null,"purl":"pkg:github/roclark/torchsr","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/roclark%2Ftorchsr","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/roclark%2Ftorchsr/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/roclark%2Ftorchsr/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/roclark%2Ftorchsr/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/roclark","download_url":"https://codeload.github.com/roclark/torchsr/tar.gz/refs/heads/main","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/roclark%2Ftorchsr/sbom","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":263986223,"owners_count":23539808,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":[],"created_at":"2024-11-21T15:07:29.792Z","updated_at":"2025-07-06T23:06:06.924Z","avatar_url":"https://github.com/roclark.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"# TorchSR\nAn easy-to-use SRGAN-based Super Resolution tool using PyTorch to upscale\nimages by 4x in each direction.\n\n## System Requirements\nThe following hardware and software versions are required to properly use this\nrepository:\n\n* Python Version: 3.7 or newer (including PIP)\n* PyTorch Version: 1.10.0 or newer (`torchrun` binary required for distributed)\n* Pillow Version: 7.1.2 or newer\n* Scikit-Learn Version: 0.18.2 or newer\n* Operating System: Linux-based (tested on Ubuntu 20.04)\n* Recommended CPU Cores: 4 or higher\n* Recommended system Memory: 16GB or higher\n* Docker Version: 19.03 or newer (if using Docker)\n\nThe application also supports Weights and Biases (WandB) for performance and\naccuracy metrics. If desired, a recent version of the `wandb` Python package can\nbe installed via `pip` to monitor these metrics.\n\n### GPU Requirements\nGPUs can dramatically improve performance by 10x or more. This repository only\nsupports recent GPUs from NVIDIA with the following requirements:\n\n* CUDA Version: 10.0 or newer\n* Driver Version: 450 or newer\n* GPU Architecture: Volta or newer\n* Recommended GPU Memory: 8GB or higher\n* NVIDIA Container Runtime (if using Docker)\n\nFor multi-GPU communication over NCCL, NVLink is highly recommended.\n\n## Setup\nTo train a new model, a dataset needs to be downloaded and extracted locally.\nWhile most any dataset can be used, it is recommended to download both DIV2K\nand FLICKR2K into a single directory for a broad dataset. TorchSR handles the\ndata preprocessing so only the high resolution images from each dataset are\nneeded for training and the downloaded low resolution images can be ignored. To\ndownload and extract the datasets, run the following steps (the download may\ntake a while depending on network speeds):\n\n```bash\n# Download the datasets\nwget http://data.vision.ee.ethz.ch/cvl/DIV2K/DIV2K_train_HR.zip\nwget http://cv.snu.ac.kr/research/EDSR/Flickr2K.tar\n# Extract both datasets\nunzip DIV2K_train_HR.zip\ntar -xvf Flickr2K.tar\n# Create a directory to save the dataset to\nmkdir -p dataset\n# Transfer all images to the dataset directory\nmv DIV2K_train_HR/*.png dataset/\nmv Flickr2K/Flickr2K_HR/*.png dataset/\n# Cleanup all archives\nrmdir DIV2K_train_HR\nrm DIV2K_train_HR.zip\nrm -rf Flickr2K\nrm Flickr2K.tar\n```\n\n**NOTE:** While most any high-resolution dataset can be used for training, it is\nrecommended to only use images saved in the PNG format as JPEGs can introduce\nartifacts on the images due to their lossy nature.\n\n### Installation\nInstall all dependencies listed in the `requirements.txt` file with `pip`:\n\n```bash\npip3 install -r requirements.txt\n```\n\nNext, to make it easier to use the application, a Python wheel can be built and\ninstalled. To build the package, run the following:\n\n```bash\npython3 setup.py bdist_wheel\n```\n\nThis will create a `dist/` directory locally which contains the newly built\nPython wheel. Install the wheel with:\n\n```bash\npip3 install dist/torchsr-*.whl\n```\n\nVerify the package was installed correctly with:\n\n```\n$ torchsr --help\nusage: torchSR Version: 0.1.0 [-h] [--local_rank N] function ...\n\npositional arguments:\n  function\n    train         Train an SRGAN model against an HD dataset.\n    test          Generated a super resolution image based on a trained SRGAN model.\n\noptional arguments:\n  -h, --help      show this help message and exit\n  --local_rank N  Rank for the local process. Used for distributed training.\n```\n\nThe remainder of the documentation assumes the Python wheel has been built and\ninstalled. If running the code directly from source and not using the wheel,\nreplace any `torchsr ...` commands below with `python3 -m torchsr.torchsr ...`\nfor identical functionality. Note that if any changes are made locally to the\ncode, the wheel will need to be rebuilt and reinstalled following the steps\nabove.\n\n### Docker\nThe repository also supports Docker containers where the package can be prebuilt\nand installed in a container with known dependencies installed. To build the\ncontainer, clone the repository and run the following inside the repository's\nroot directory:\n\n```\ndocker build -t torchsr:0.1.0 .\n```\n\nThis will build a new Docker container with the latest requirements and code.\n\nTo launch the container, run the following:\n\n```\ndocker run --rm -it \\\n    --gpus all \\  # If using GPUs\n    -v /path/to/dataset:/dataset \\\n    --shm-size=16g \\\n    torchsr:0.1.0\n```\n\nUpdate the `-v /path/to/dataset:/dataset` to the path to where the dataset was\ndownloaded during the setup. This will map the directory inside the container\nwhere it can be used freely. Omit the `--gpus all` flag if not using GPUs.\n\nOnce inside the container, the `torchsr` command is in the PATH and all commands\nbelow will run by default.\n\n## Running\nWith the application fully configured and installed, training can begin. TorchSR\nattempts to automatically set as many parameters as possible to make the\napplication easy to use. To use the default parameters, simply run\n\n```bash\ntorchsr train\n```\n\nto initiate the training process. By default, this will use the dataset saved to\n`dataset/` locally and run on as many supported GPUs as possible.\n\n### Advanced Usage\nCertain parameters can be used to further tune performance and accuracy\ndepending on the dataset and local environment among other variables. Here are\nsome of the more common parameters that can be specified for different\nfunctionality:\n\n* `--batch-size N`: Use a different batch size from the default. While any\npositive integer can be used, it is recommended to use a number that is a power\nof 2 (2, 4, 8, 16, ...). In general, the higher the batch size, the faster the\ntraining will run with diminishing returns while also increasing the amount of\nmemory allocated by the GPUs. If you run into memory issues, whether CPU or GPU\nmemory, try dropping the batch size in half and re-running.\n* `--dataset-multiplier N`: TorchSR takes a single random 96x96 cropping of each\nimage in the training dataset in every epoch and downscales the cropping to\ncreate a low and high resolution source. This process wastes over 99% of the\nimage in the dataset as everything outside of the 96x96 cropping is thrown out\nfor each epoch. Using the `--dataset-multiplier` flag increases the number of\nrandom samples from each image, growing the total dataset size without\nincreasing on-disk requirements. The number for the multiplier is the number of\nrandom samples to take from every image.\n* `--train-dir \u003cdirectory\u003e`: If the dataset is saved in a directory other than\n`dataset/`, specify the location by passing it to `--train-dir`.\n\n### Weights and Biases\nSupport for Weights and Biases (WandB) has been included with the repository to\nmake it easier to view training progress and share results. To get started,\nsignup for a free account on WandB's website and copy your API key from your\ndashboard. To have WandB push your results to its web service automatically,\nsimply install the `wandb` Python package and login to WandB from the CLI:\n\n```\n$ pip3 install wandb\n$ wandb login\nwandb: You can find your API key in your browser here: https://wandb.ai/authorize\nwandb: Paste an API key from your profile and hit enter, or press ctrl+c to quit: \u003cEnter API Key Here\u003e\n```\n\nOnce logged in, run TorchSR following any of the listed commands on this page.\nTraining logs can be viewed at https://wandb.ai.\n\n### Distributed Training\nTorchSR can be run on multiple GPUs to greatly accelerate training by using the\n`torchrun` binary included with modern versions of PyTorch. To launch a\nmulti-GPU training pass, modify how the application is called by adding\n`torchrun --nproc_per_node=X` before your command. For example:\n\n```bash\n# Original 1-GPU command:\ntorchsr train --dataset-multiplier 10\n\n# Equivalent command running on 8 GPUs on a single node:\ntorchrun --nproc_per_node=8 -m torchsr.torchsr train --dataset-multiplier 10\n```\n\nWhile running distributed training, the dataset will be divided between each\nGPU, reducing the overall time needed to complete a single epoch.\n\n### Slurm (GPU only)\nTorchSR supports running on Slurm clusters using [Pyxis](https://github.com/NVIDIA/pyxis)\nand [Enroot](https://github.com/NVIDIA/enroot). To run on a Slurm cluster with\nPyxis and Enroot, build the Docker image following the steps above and push the\nimage to a container registry of your choice using\n`docker push \u003cimage name:tag\u003e`. To launch a training job on Slurm, create an\nsbatch file named `distributed.sh` with the following contents:\n\n```bash\n#!/bin/bash\n#SBATCH --exclusive\n#SBATCH --ntasks-per-node=8\n\nsrun \\\n    --gpus-per-node=8 \\\n    --container-image \u003cimage name:tag here\u003e \\\n    --container-mounts=/path/to/dataset:/dataset \\\n    torchsr train \\\n        --train-dir /dataset \\\n        --dataset-multiplier 10\n```\n\nModify the `--ntasks-per-node=X` and `--gpus-per-node=X` values to match the expected\nnumber of GPUs available per node. Replace `\u003cimage name:tag here\u003e` with the name\nspecified during the Docker image build above. Replace the `/path/to/dataset`\nwith the location of your dataset.\n\nTo launch the Slurm job, run `sbatch --nodes=X distributed.sh` and specify the\nnumber of nodes to run on with the `--nodes=X` flag. This will create a file in\nthe local directory named `slurm-XXXX.out` where `XXXX` is the job ID. This file\ncontains information related to the training job.\n\nIf using WandB, prior to launching the job, specify your WandB API key with an\nenvironment variable:\n\n```\nexport WANDB_API_KEY=\u003cEnter your key here\u003e\n```\n\nNext, modify the `distributed.sh` script to the following:\n\n```bash\n#!/bin/bash\n#SBATCH --exclusive\n#SBATCH --ntasks-per-node=8\n\nsrun \\\n    --gpus-per-node=8 \\\n    --container-image \u003cimage name:tag here\u003e \\\n    --container-mounts=/path/to/dataset:/dataset \\\n    bash -c 'wandb login \u0026\u0026 \\\n    torchsr train \\\n        --train-dir /dataset \\\n        --dataset-multiplier 10'\n```\n\nOnce modified, the job can be launched as before using `sbatch` and WandB will\nbe authenticated and updated automatically.\n\n### Monitoring\nTo help view progress while training, a validation image is generated based on\nthe latest model weights after every epoch. The image is saved to the `output`\ndirectory with the epoch number included in every filename. The image is a\nroyalty free image from Pixabay depicting a waterfall in nature. There are a lot\nof colors and details in this model which allow for a good way to showcase model\naccuracy. The images below are of the low resolution input on the left and the\nupscaled image that was generated on the right.\n\n| Low Resolution                                                    | Super Resolution                                        |\n|:------------------------------------------------------------------|:-------------------------------------------------|\n| \u003cimg src=\"media/waterfalls-low-res.png\" width=\"480\" height=\"320\"\u003e | \u003cimg src=\"media/waterfalls-super-res.png\" width=\"480\" height=\"320\"\u003e |\n\n## References\nThis project is based on the research published in the following paper:\n**Photo-Realistic Single Image Super-Resolution Using a Generative Adversarial Network**\n\n```\n@InProceedings{srgan,\n    author = {Christian Ledig, Lucas Theis, Ferenc Huszar, Jose Caballero, Andrew Cunningham, Alejandro Acosta, Andrew Aitken, Alykhan Tejani, Johannes Totz, Zehan Wang, Wenzhe Shi},\n    title = {Photo-Realistic Single Image Super-Resolution Using a Generative Adversarial Network},\n    booktitle = {arXiv},\n    year = {2016}\n}\n```\n\nAdditionally, this project leverages resources from a few open-source\nrepositories:\n\n* SRGAN-PyTorch: https://github.com/Lornatang/SRGAN-PyTorch\n* BasicSR: https://github.com/xinntao/BasicSR\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Froclark%2Ftorchsr","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Froclark%2Ftorchsr","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Froclark%2Ftorchsr/lists"}