Update CI adding Python 3.9, 3.10, 3.11, and 3.12

.github/workflows/download-emnist.yml

@@ -0,0 +1,33 @@
+name: Download emnist
+
+on: [push]
+
+jobs:
+  download:
+
+    runs-on: ubuntu-latest
+    strategy:
+      max-parallel: 8
+      matrix:
+        torchvision-version: [0.5.0, 0.6.0, 0.6.1, 0.7.0, 0.8.0, 0.8.1, 0.8.2, 0.9.0, 0.9.1, 0.10.0, 0.10.1, 0.11.1, 0.11.2, 0.11.3, 0.12.0, 0.13.0, 0.13.1, 0.14.0, 0.14.1, 0.15.1, 0.15.2, 0.16.0, 0.16.1, 0.16.2, 0.17.0, 0.17.1, 0.17.2, 0.18.0]
+
+    steps:
+    - uses: actions/checkout@v4
+    - name: Set up Python 3.8
+      uses: actions/setup-python@v5
+      with:
+        python-version: 3.8
+    - name: Install dependencies
+      run: |
+        sudo apt-get update && sudo apt-get install -y libtinfo5
+        python -m pip install --upgrade pip
+        pip install numpy -U
+        pip install torchvision==${{ matrix.torchvision-version }} -f https://download.pytorch.org/whl/torch_stable.html
+        pip install setuptools
+        pip install requests
+    - name: Install package
+      run: python -m pip install .
+    - name: Download EMNIST dataset
+      run: python examples/emnist/download.py
+    - name: Extract EMNIST dataset
+      run: python examples/emnist/main.py --epochs 0

.github/workflows/example-emnist.yml

@@ -0,0 +1,51 @@
+name: Example emnist
+
+on: [push]
+
+jobs:
+  train:
+
+    runs-on: ${{ matrix.os }}
+    strategy:
+      max-parallel: 8
+      matrix:
+        python-version: [3.7, 3.8, 3.9, '3.10', 3.11, 3.12]
+        os: [macos-13, windows-latest, ubuntu-latest]
+        include:
+          - pytorch-version: 1.4.0
+            torchvision-version: 0.5.0
+          - pytorch-version: 1.8.0
+            torchvision-version: 0.9.0
+            python-version: 3.9
+          - pytorch-version: 1.11.0
+            torchvision-version: 0.12.0
+            python-version: '3.10'
+          - pytorch-version: 2.0.0
+            torchvision-version: 0.15.1
+            python-version: 3.11
+          - pytorch-version: 2.2.0
+            torchvision-version: 0.17.0 
+            python-version: 3.12
+          - pytorch-option: '+cpu'
+          - pytorch-option: ''
+            os: macos-13
+
+    steps:
+    - uses: actions/checkout@v4
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@v5
+      with:
+        python-version: ${{ matrix.python-version }}
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install numpy -U
+        pip install torch==${{ matrix.pytorch-version }}${{ matrix.pytorch-option }} torchvision==${{ matrix.torchvision-version }}${{ matrix.pytorch-option }} -f https://download.pytorch.org/whl/torch_stable.html
+        pip install setuptools
+        pip install requests
+    - name: Install package
+      run: python -m pip install .
+    - name: Download EMNIST dataset
+      run: python examples/emnist/download.py
+    - name: Train a model on EMNIST dataset
+      run: python examples/emnist/main.py --epochs 1

.github/workflows/example-plots.yml

@@ -3,14 +3,27 @@ name: Example plots
 on: [push]
 
 jobs:
-  build:
+  plot:
 
     runs-on: ${{ matrix.os }}
     strategy:
       max-parallel: 8
       matrix:
-        python-version: [3.7, 3.8]
+        python-version: [3.7, 3.8, 3.9, '3.10', 3.11, 3.12]
         os: [macos-13, windows-latest, ubuntu-latest]
+        include:
+          - pytorch-version: 1.4.0
+          - pytorch-version: 1.8.0
+            python-version: 3.9
+          - pytorch-version: 1.11.0
+            python-version: '3.10'
+          - pytorch-version: 2.0.0
+            python-version: 3.11
+          - pytorch-version: 2.2.0
+            python-version: 3.12
+          - pytorch-option: '+cpu'
+          - pytorch-option: ''
+            os: macos-13
 
     steps:
     - uses: actions/checkout@v4
@@ -22,10 +35,11 @@ jobs:
       run: |
         python -m pip install --upgrade pip
         pip install numpy -U
-        pip install torch==1.4.0 -f https://download.pytorch.org/whl/torch_stable.html
+        pip install torch==${{ matrix.pytorch-version }}${{ matrix.pytorch-option }} -f https://download.pytorch.org/whl/torch_stable.html
         pip install matplotlib
+        pip install setuptools
     - name: Install package
-      run: python setup.py install
+      run: python -m pip install .
     - name: Preparation
       run: mkdir artifact
     - name: Plot warmup period

.github/workflows/pythonpackage.yml

@@ -9,8 +9,21 @@ jobs:
     strategy:
       max-parallel: 8
       matrix:
-        python-version: [3.7, 3.8]
+        python-version: [3.7, 3.8, 3.9, '3.10', 3.11, 3.12]
         os: [macos-13, windows-latest, ubuntu-latest]
+        include:
+          - pytorch-version: 1.4.0
+          - pytorch-version: 1.8.0
+            python-version: 3.9
+          - pytorch-version: 1.11.0
+            python-version: '3.10'
+          - pytorch-version: 2.0.0
+            python-version: 3.11
+          - pytorch-version: 2.2.0
+            python-version: 3.12
+          - pytorch-option: '+cpu'
+          - pytorch-option: ''
+            os: macos-13
 
     steps:
     - uses: actions/checkout@v4
@@ -22,7 +35,7 @@ jobs:
       run: |
         python -m pip install --upgrade pip
         pip install numpy -U
-        pip install torch==1.4.0 -f https://download.pytorch.org/whl/torch_stable.html
+        pip install torch==${{ matrix.pytorch-version }}${{ matrix.pytorch-option }} -f https://download.pytorch.org/whl/torch_stable.html
     - name: Lint with flake8
       run: |
         pip install flake8
@@ -34,3 +47,7 @@ jobs:
       run: |
         pip install pytest
         pytest test -s -vv
+    - name: Package with build
+      run: |
+        pip install setuptools build
+        python -m build

README.md

@@ -11,13 +11,13 @@ This library contains PyTorch implementations of the warmup schedules described
 
 ## Installation
 
-Make sure you have Python 3.6+ and PyTorch 1.1+. Then, run the following command:
+Make sure you have Python 3.7+ and PyTorch 1.1+. Then, run the following command in the project directory:
 
 ```
-python setup.py install
+python -m pip install .
 ```
 
-or
+or install the latest version from the Python Package Index:
 
 ```
 pip install -U pytorch_warmup
@@ -92,12 +92,12 @@ When the learning rate schedule uses the epoch number, the warmup schedule can b
 lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[num_epochs//3], gamma=0.1)
 warmup_scheduler = warmup.UntunedLinearWarmup(optimizer)
 for epoch in range(1,num_epochs+1):
-    for iter, batch in enumerate(dataloader):
+    for i, batch in enumerate(dataloader):
         optimizer.zero_grad()
         loss = ...
         loss.backward()
         optimizer.step()
-        if iter < len(dataloader)-1:
+        if i < len(dataloader)-1:
             with warmup_scheduler.dampening():
                 pass
     with warmup_scheduler.dampening():
@@ -108,16 +108,36 @@ This code can be rewritten more compactly:
 
 ```python
 for epoch in range(1,num_epochs+1):
-    for iter, batch in enumerate(dataloader):
+    for i, batch in enumerate(dataloader):
         optimizer.zero_grad()
         loss = ...
         loss.backward()
         optimizer.step()
         with warmup_scheduler.dampening():
-            if iter + 1 == len(dataloader):
+            if i + 1 == len(dataloader):
                 lr_scheduler.step()
 ```
 
+#### Approach 3
+When you use `CosineAnnealingWarmRestarts`, the warmup schedule can be used as follows:
+
+```python
+lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0=10, T_mult=2)
+warmup_period = 2000
+warmup_scheduler = warmup.LinearWarmup(optimizer, warmup_period)
+iters = len(dataloader)
+warmup_epochs = ... # for example, (warmup_period + iters - 1) // iters
+for epoch in range(epochs+warmup_epochs):
+    for i, batch in enumerate(dataloader):
+        optimizer.zero_grad()
+        loss = ...
+        loss.backward()
+        optimizer.step()
+        with warmup_scheduler.dampening():
+            if epoch >= warmup_epochs:
+                lr_scheduler.step(epoch-warmup_epochs + i / iters)
+```
+
 ### Warmup Schedules
 
 #### Manual Warmup

examples/plots/effective_warmup_period.py

@@ -1,6 +1,6 @@
 import argparse
 import matplotlib.pyplot as plt
-from pytorch_warmup import *
+from pytorch_warmup import rho_fn, rho_inf_fn, get_offset
 import numpy as np
 
 

examples/plots/warmup_schedule.py

@@ -1,7 +1,7 @@
 import argparse
 import matplotlib.pyplot as plt
 import torch
-from pytorch_warmup import *
+from pytorch_warmup import RAdamWarmup, UntunedExponentialWarmup, UntunedLinearWarmup
 
 
 def get_rates(warmup_cls, beta2, max_step):

pytorch_warmup/__init__.py

@@ -1,3 +1,17 @@
-from .base import *
-from .untuned import *
-from .radam import *
+from .base import BaseWarmup, LinearWarmup, ExponentialWarmup
+from .untuned import UntunedLinearWarmup, UntunedExponentialWarmup
+from .radam import RAdamWarmup, rho_fn, rho_inf_fn, get_offset
+
+__version__ = "0.2.0.dev0"
+
+__all__ = [
+    'BaseWarmup',
+    'LinearWarmup',
+    'ExponentialWarmup',
+    'UntunedLinearWarmup',
+    'UntunedExponentialWarmup',
+    'RAdamWarmup',
+    'rho_fn',
+    'rho_inf_fn',
+    'get_offset',
+]

pytorch_warmup/base.py

@@ -66,17 +66,17 @@ def warmup_factor(self, step, **params):
 
 
 def get_warmup_params(warmup_period, group_count):
-    if type(warmup_period) == list:
+    if isinstance(warmup_period, list):
         if len(warmup_period) != group_count:
             raise ValueError(
                 'size of warmup_period does not equal {}.'.format(group_count))
         for x in warmup_period:
-            if type(x) != int:
+            if not isinstance(x, int):
                 raise ValueError(
                     'An element in warmup_period, {}, is not an int.'.format(
                         type(x).__name__))
         warmup_params = [dict(warmup_period=x) for x in warmup_period]
-    elif type(warmup_period) == int:
+    elif isinstance(warmup_period, int):
         warmup_params = [dict(warmup_period=warmup_period)
                          for _ in range(group_count)]
     else:

setup.cfg

@@ -1,2 +1,3 @@
 [metadata]
+version = attr: pytorch_warmup.__version__
 license_files = LICENSE

setup.py

@@ -5,7 +5,6 @@
 
 setuptools.setup(
     name="pytorch-warmup",
-    version="0.1.1",
     author="Takenori Yamamoto",
     author_email="yamamoto.takenory@gmail.com",
     description="A PyTorch Extension for Learning Rate Warmup",
@@ -14,11 +13,22 @@
     url="https://github.com/Tony-Y/pytorch_warmup",
     packages=['pytorch_warmup'],
     classifiers=[
+        "Development Status :: 5 - Production/Stable",
+        "Intended Audience :: Developers",
+        "Intended Audience :: Education",
         "Intended Audience :: Science/Research",
+        "Topic :: Scientific/Engineering",
         "Topic :: Scientific/Engineering :: Artificial Intelligence",
+        "Topic :: Software Development",
+        "Topic :: Software Development :: Libraries",
         "Topic :: Software Development :: Libraries :: Python Modules",
+        "Programming Language :: Python :: 3",
         "Programming Language :: Python :: 3.7",
         "Programming Language :: Python :: 3.8",
+        "Programming Language :: Python :: 3.9",
+        "Programming Language :: Python :: 3.10",
+        "Programming Language :: Python :: 3.11",
+        "Programming Language :: Python :: 3.12",
         "Operating System :: OS Independent",
         "License :: OSI Approved :: MIT License",
     ],