Code covarge and setup.py

Author: BlackHC

.coveragerc

@@ -0,0 +1,5 @@
+[run]
+source=toma/*
+omit=
+    */tests/*
+    setup.py

README.md

@@ -1,7 +1,13 @@
-# TOrch Memory-Adaptive Algorithms
+# Torch Memory-adaptive Algorithms (TOMA)
+
+[![Build Status](https://www.travis-ci.com/BlackHC/implicit_lambda.svg?branch=master)](https://www.travis-ci.com/BlackHC/implicit_lambda) [![codecov](https://codecov.io/gh/BlackHC/implicit_lambda/branch/master/graph/badge.svg)](https://codecov.io/gh/BlackHC/implicit_lambda) [![PyPI](https://img.shields.io/badge/PyPI-implicit_lambda-blue.svg)](https://pypi.python.org/pypi/implicit_lambda/)
+
+
+A collection of helpers to make it easier to write code that adapts to the available CUDA memory.
+
+## Installation
+
 
 TODOs:
 * [] write readme
-* [] add a test for simple_chunked
-* [x] add tests for explicit_*
-* [] add tests for the decorators
+* [] add doc strings

setup.cfg

@@ -2,7 +2,7 @@
 test=pytest
 
 [tool:pytest]
-addopts = --cov=.
+addopts = --cov toma
 
 [pylama:pycodestyle]
 max_line_length = 120

setup.py

@@ -0,0 +1,65 @@
+# Always prefer setuptools over distutils
+from setuptools import setup
+
+# To use a consistent encoding
+from codecs import open
+from os import path
+
+here = path.abspath(path.dirname(__file__))
+
+# Get the long description from the README file
+with open(path.join(here, "README.md"), encoding="utf-8") as f:
+    long_description = f.read()
+
+setup(
+    name="toma",
+    # Versions should comply with PEP440.  For a discussion on single-sourcing
+    # the version across setup.py and the project code, see
+    # https://packaging.python.org/en/latest/single_source_version.html
+    version="0.0.0",
+    description="Make it easy to write algorithms in PyTorch that adapt to the available CUDA memory",
+    # Fix windows newlines.
+    long_description=long_description.replace("\r\n", "\n"),
+    # The project's main homepage.
+    url="https://github.com/blackhc/toma",
+    # Author details
+    author="Andreas @blackhc Kirsch",
+    author_email="[email protected]",
+    # Choose your license
+    license="MIT",
+    # See https://pypi.python.org/pypi?%3Aaction=list_classifiers
+    classifiers=[
+        # How mature is this project? Common values are
+        #   3 - Alpha
+        #   4 - Beta
+        #   5 - Production/Stable
+        "Development Status :: 3 - Alpha",
+        # Indicate who your project is intended for
+        "Intended Audience :: Developers",
+        "Intended Audience :: Science/Research",
+        "Topic :: Software Development :: Libraries :: Python Modules",
+        # Pick your license as you wish (should match "license" above)
+        "License :: OSI Approved :: MIT License",
+        "Programming Language :: Python :: 3.7",
+    ],
+    # What does your project relate to?
+    keywords="tools pytorch",
+    # You can just specify the packages manually here if your project is
+    # simple. Or you can use find_packages().
+    packages=["toma"],
+    package_dir={"": "src"},
+    # List run-time dependencies here.  These will be installed by pip when
+    # your project is installed. For an analysis of "install_requires" vs pip's
+    # requirements files see:
+    # https://packaging.python.org/en/latest/requirements.html
+    install_requires=[],
+    # List additional groups of dependencies here (e.g. development
+    # dependencies). You can install these using the following syntax,
+    # for example:
+    # $ pip install -e .[dev,test]
+    extras_require={
+        "dev": ["check-manifest"],
+        "test": ["coverage", "codecov", "pytest", "pytest-benchmark", "pytest-cov"],
+    },
+    setup_requires=["pytest-runner"],
+)

tests/test_explicit_toma.py

@@ -1,3 +1,6 @@
+import torch
+import pytest
+
 from toma import explicit
 from toma.batchsize_cache import NoBatchsizeCache, GlobalBatchsizeCache, StacktraceMemoryBatchsizeCache
 
@@ -6,7 +9,23 @@ def raise_fake_oom():
     raise RuntimeError("CUDA out of memory.")
 
 
-def test_fake_explicit_toma_none():
+def test_fake_explicit_batch_raise():
+    def f(batchsize):
+        raise_fake_oom()
+
+    with pytest.raises(RuntimeError):
+        explicit.batch(f, 64)
+
+
+def test_fake_explicit_range_raise():
+    def f(start, end):
+        raise_fake_oom()
+
+    with pytest.raises(RuntimeError):
+        explicit.range(f, 0, 64, 64)
+
+
+def test_fake_explicit_batch_none():
     batchsizes = []
 
     def f(batchsize):
@@ -22,7 +41,7 @@ def f(batchsize):
     assert batchsizes == [64, 32, 16, 64, 32, 16]
 
 
-def test_fake_explicit_toma_global():
+def test_fake_explicit_batch_global():
     batchsizes = []
 
     def f(batchsize):
@@ -40,7 +59,7 @@ def f(batchsize):
     assert batchsizes == [64, 32, 16, 16, 16, 16]
 
 
-def test_fake_explicit_toma_sm():
+def test_fake_explicit_batch_sm():
     batchsizes = []
 
     def f(batchsize):
@@ -58,7 +77,7 @@ def f(batchsize):
     assert batchsizes == [64, 32, 16, 16, 16, 64, 32, 16]
 
 
-def test_fake_explicit_toma_mix():
+def test_fake_explicit_batch_mix():
     batchsizes = []
 
     def f(batchsize):
@@ -79,7 +98,7 @@ def f(batchsize):
     assert batchsizes == [64, 32, 16, 16, 16, 16, 64, 32, 16, 16]
 
 
-def test_fake_explicit_toma_range_none():
+def test_fake_explicit_range_none():
     batchsizes = []
 
     def f(start, end):
@@ -99,7 +118,7 @@ def f(start, end):
     assert batchsizes == [64, 64, 32, 32, 16, 16] * 2
 
 
-def test_fake_explicit_toma_range_global():
+def test_fake_explicit_range_global():
     batchsizes = []
 
     def f(start, end):
@@ -121,7 +140,7 @@ def f(start, end):
     assert batchsizes == [64, 64, 32, 32, 16, 16] + [16] * 8 * 2
 
 
-def test_fake_explicit_toma_range_sm():
+def test_fake_explicit_range_sm():
     batchsizes = []
 
     def f(start, end):
@@ -143,7 +162,7 @@ def f(start, end):
     assert batchsizes == [64, 64, 32, 32, 16, 16] + [16] * 8 + [64, 64, 32, 32, 16, 16]
 
 
-def test_fake_explicit_toma_range_sm():
+def test_fake_explicit_range_sm():
     batchsizes = []
 
     def f(start, end):
@@ -166,3 +185,12 @@ def f(start, end):
         explicit.range(f, 0, 128, 64, toma_cache_type=StacktraceMemoryBatchsizeCache)
 
     assert batchsizes == ([64, 64, 32, 32, 16, 16] + [16] * 8 * 2 + [64, 64, 32, 32, 16, 16] + [16] * 8)
+
+
+def test_explicit_chunked():
+    def func(tensor, start, end):
+        tensor[:] = 1.
+
+    tensor = torch.zeros((128, 4, 4))
+    explicit.chunked(func, tensor, 32)
+    assert torch.allclose(tensor, torch.tensor(1.))

tests/test_simple_toma.py

@@ -7,7 +7,7 @@ def raise_fake_oom():
     raise RuntimeError("CUDA out of memory.")
 
 
-def test_fake_simple_toma():
+def test_fake_simple_batch():
     hit_16 = False
 
     def f(batch_size):
@@ -24,7 +24,7 @@ def f(batch_size):
     assert hit_16
 
 
-def test_fake_simple_toma_range():
+def test_fake_simple_range():
     hit_16 = False
 
     def f(start, end):
@@ -43,7 +43,7 @@ def f(start, end):
     assert hit_16
 
 
-def test_fake_simple_toma_chunked():
+def test_fake_simple_chunked():
     hit_16 = False
 
     def f(tensor, start, end):
@@ -63,7 +63,7 @@ def f(tensor, start, end):
     assert hit_16
 
 
-def test_simple_toma():
+def test_simple_batch():
     import torch
 
     if not torch.cuda.is_available():
@@ -89,7 +89,7 @@ def f(batch_size):
     assert succeeded
 
 
-def test_simple_toma_range():
+def test_simple_range():
     import torch
 
     if not torch.cuda.is_available():

tests/test_toma.py

@@ -1,11 +1,12 @@
+import torch
 from toma import toma, explicit, batchsize_cache as tbc
 
 
 def raise_fake_oom():
     raise RuntimeError("CUDA out of memory.")
 
 
-def test_fake_toma_simple():
+def test_fake_batch_none():
     batchsizes = []
 
     @toma.batch(initial_batchsize=64, cache_type=tbc.NoBatchsizeCache)
@@ -22,7 +23,7 @@ def f(batchsize):
     assert batchsizes == [64, 32, 16, 64, 32, 16]
 
 
-def test_fake_toma_explicit():
+def test_fake_batch_global():
     batchsizes = []
 
     @toma.batch(initial_batchsize=64, cache_type=tbc.GlobalBatchsizeCache)
@@ -39,7 +40,7 @@ def f(batchsize):
     assert batchsizes == [64, 32, 16, 16]
 
 
-def test_fake_toma_range_global():
+def test_fake_range_none():
     batchsizes = []
 
     @toma.range(initial_step=64, cache_type=tbc.NoBatchsizeCache)
@@ -60,7 +61,7 @@ def f(start, end):
     assert batchsizes == [64, 64, 32, 32, 16, 16] * 2
 
 
-def test_fake_toma_range_explicit():
+def test_fake_range_global():
     batchsizes = []
 
     @toma.range(initial_step=64, cache_type=tbc.GlobalBatchsizeCache)
@@ -79,3 +80,13 @@ def f(start, end):
         f(0, 128)
 
     assert batchsizes == [64, 64, 32, 32, 16, 16] + [16] * 8
+
+
+def test_chunked():
+    @toma.chunked(initial_step=32)
+    def func(tensor, start, end):
+        tensor[:] = 1.
+
+    tensor = torch.zeros((128, 4, 4))
+    func(tensor)
+    assert torch.allclose(tensor, torch.tensor(1.))

src/toma/__init__.py toma/__init__.pysrc/toma/__init__.py renamed to toma/__init__.py

@@ -178,7 +178,7 @@ def batch(
 
         while True:
             try:
-                value = batchsize.get_batchsize()
+                value = batchsize.get()
                 return func(value, *args, **kwargs)
             except RuntimeError as exception:
                 if value > 1 and tcm.should_reduce_batch_size(exception):
@@ -206,10 +206,10 @@ def range(
         current = start
         while current < end:
             try:
-                func(current, min(current + batchsize.get_batchsize(), end), *args, **kwargs)
-                current += batchsize.get_batchsize()
+                func(current, min(current + batchsize.get(), end), *args, **kwargs)
+                current += batchsize.get()
             except RuntimeError as exception:
-                if batchsize.get_batchsize() > 1 and tcm.should_reduce_batch_size(exception):
+                if batchsize.get() > 1 and tcm.should_reduce_batch_size(exception):
                     batchsize.decrease_batchsize()
                     tcm.gc_cuda()
                 else:

src/toma/batchsize_cache.py toma/batchsize_cache.pysrc/toma/batchsize_cache.py renamed to toma/batchsize_cache.py

@@ -14,11 +14,12 @@ def set_initial_batchsize(self, initial_batchsize: int):
         if not self.value:
             self.value = initial_batchsize
 
-    def get_batchsize(self) -> int:
+    def get(self) -> int:
         return self.value
 
     def decrease_batchsize(self):
         self.value //= 2
+        assert self.value > 0
 
 
 class BatchsizeCache:

src/toma/stacktrace.py toma/stacktrace.pysrc/toma/stacktrace.py renamed to toma/stacktrace.py

@@ -16,7 +16,7 @@ def _constant_code_context(code_context):
 def get_simple_traceback(ignore_top=0):
     """Get a simple trackback that can be hashed and won't create reference
     cyles."""
-    stack = inspect.stack(context=ignore_top + 1)[ignore_top : -__watermark - 1]
+    stack = inspect.stack(context=1)[ignore_top + 1: -__watermark - 1]
     simple_traceback = tuple(
         (fi.filename, fi.lineno, fi.function, _constant_code_context(fi.code_context), fi.index) for fi in stack
     )