Is albumentations support useful? about ffcv

[GuillaumeLeclerc]

Hi! Thanks for the suggestion! We are of course interested in as many augmentations as possible so that sounds like a great addition. Your commit looks very reasonable to me and should be a good start! I'm wondering if it would be possible for us to make it even faster by integrating with the JIT pipeline. Can you give me some details on how they are actually implemented under the hood so that I can estimate feasibility.

[ar90n]

Thanks for your quick reply! I evaluated this feature by modifying the pipeline of train_cifar.py such as follows.

@@ -45,6 +45,8 @@ from ffcv.transforms import RandomHorizontalFlip, Cutout, \
 from ffcv.transforms.common import Squeeze
 from ffcv.writer import DatasetWriter

+import albumentations as A
+
 Section('training', 'Hyperparameters').params(
     lr=Param(float, 'The learning rate to use', required=True),
     epochs=Param(int, 'Number of epochs to run for', required=True),
@@ -83,9 +85,14 @@ def make_dataloaders(train_dataset=None, val_dataset=None, batch_size=None, num_
         image_pipeline: List[Operation] = [SimpleRGBImageDecoder()]
         if name == 'train':
             image_pipeline.extend([
-                RandomHorizontalFlip(),
-                RandomTranslate(padding=2, fill=tuple(map(int, CIFAR_MEAN))),
-                Cutout(4, tuple(map(int, CIFAR_MEAN))),
+                A.HorizontalFlip(),
+                A.ShiftScaleRotate(scale_limit=0.0, rotate_limit=0.0),
+                A.Cutout(
+                    num_holes=1,
+                    max_h_size=4,
+                    max_w_size=4,
+                    fill_value=tuple(map(int, CIFAR_MEAN)),
+                )
             ])
         image_pipeline.extend([
             ToTensor(),

The above codes are almost equivalent to the original codes. And it worked correctly.
And I think it is useful that we can use jit compiled operations and wrapped albumations operations in the same pipeline.
So I evaluated the following codes. These codes are also almost equivalent to original codes.

@@ -45,6 +45,8 @@ from ffcv.transforms import RandomHorizontalFlip, Cutout, \
 from ffcv.transforms.common import Squeeze
 from ffcv.writer import DatasetWriter

+import albumentations as A
+
 Section('training', 'Hyperparameters').params(
     lr=Param(float, 'The learning rate to use', required=True),
     epochs=Param(int, 'Number of epochs to run for', required=True),
@@ -83,7 +85,7 @@ def make_dataloaders(train_dataset=None, val_dataset=None, batch_size=None, num_
         image_pipeline: List[Operation] = [SimpleRGBImageDecoder()]
         if name == 'train':
             image_pipeline.extend([
-                RandomHorizontalFlip(),
+                A.HorizontalFlip(),
                 RandomTranslate(padding=2, fill=tuple(map(int, CIFAR_MEAN))),
                 Cutout(4, tuple(map(int, CIFAR_MEAN))),
             ])

But it didn't work. I met the following assertion error.

Traceback (most recent call last):
  File "/workspaces/ffcv/examples/cifar/train_cifar.py", line 210, in <module>
    loaders, start_time = make_dataloaders()
  File "/root/.local/lib/python3.9/site-packages/fastargs/decorators.py", line 63, in result
    return func(*args, **kwargs)
  File "/root/.local/lib/python3.9/site-packages/fastargs/decorators.py", line 35, in __call__
    return self.func(*args, **filled_args)
  File "/workspaces/ffcv/examples/cifar/train_cifar.py", line 102, in make_dataloaders
    loaders[name] = Loader(paths[name], batch_size=batch_size, num_workers=num_workers,
  File "/workspaces/ffcv/ffcv/loader/loader.py", line 196, in __init__
    self.pipelines[field_name] = Pipeline(operations)
  File "/workspaces/ffcv/ffcv/pipeline/pipeline.py", line 25, in __init__
    self.operation_blocks, _ = self.parse_pipeline()
  File "/workspaces/ffcv/ffcv/pipeline/pipeline.py", line 42, in parse_pipeline
    current_state, memory_allocation = operation.declare_state_and_memory(
  File "/workspaces/ffcv/ffcv/transforms/translate.py", line 54, in declare_state_and_memory
    assert previous_state.jit_mode
AssertionError

It seems that such operations which have assert previous_state.jit_mode in their declare_state_and_memory require that previous operations must be jit-compiled.
Is my understanding correct? Are there any solutions?

[GuillaumeLeclerc]

I'm not sure about this assertion. @andrewilyas any idea why it is required.

Either way to integrate them propery in the JIT pipeline I would need to know how are the augmentations implemented internally, does it use a python extension or something similar ?

[andrewilyas]

Thanks for pointing this out @ar90n ! Those assertions can be removed and replaced with a simple setter, most likely. I'll start working on a fix and keep you updated!

[ar90n]

@andrewilyas
I tried your PR(#59) and got the following errors.

...

  0%|                                                                                                                                                                                                                                                                                                    | 0/97 [00:00<?, ?it/s]Exception in thread Thread-2:
Traceback (most recent call last):
  File "/opt/miniconda/envs/ffcv/lib/python3.9/threading.py", line 973, in _bootstrap_inner
    self.run()
  File "/workspaces/ffcv/ffcv/loader/epoch_iterator.py", line 80, in run
    result = self.run_pipeline(b_ix, ixes, slot, events[slot])
  File "/workspaces/ffcv/ffcv/loader/epoch_iterator.py", line 134, in run_pipeline
    result = code(*args)
  File "/opt/miniconda/envs/ffcv/lib/python3.9/site-packages/numba/core/dispatcher.py", line 468, in _compile_for_args
    error_rewrite(e, 'typing')
  File "/opt/miniconda/envs/ffcv/lib/python3.9/site-packages/numba/core/dispatcher.py", line 409, in error_rewrite
    raise e.with_traceback(None)
numba.core.errors.TypingError: Failed in nopython mode pipeline (step: nopython frontend)
non-precise type pyobject
During: typing of argument at  (2)

I think this error has occurred when jit compiled operations run after non-jit compiled operations.
To evaluate this assumption, I wrote the following simple Buffer operation.

class Buffer(Operation):

    def __init__(self, jit_mode):
        super().__init__()
        self.jit_mode = jit_mode

    def generate_code(self) -> Callable:
        my_range = Compiler.get_iterator()

        def apply_transform(images, dst):
            for i in my_range(images.shape[0]):
                dst[i] = images[i]
            return dst

        return apply_transform

    def declare_state_and_memory(
        self, previous_state: State
    ) -> Tuple[State, Optional[AllocationQuery]]:
        return (replace(previous_state, jit_mode=self.jit_mode), \
                AllocationQuery(shape=previous_state.shape, dtype=previous_state.dtype))

And I evaluated the following simple tests. These tests are motivated to reveal the reason of above error.

• jit operation -> non-jit operation

       image_pipeline: List[Operation] = [SimpleRGBImageDecoder()]
        if name == 'train':
            image_pipeline.extend([
                Buffer(True),
                Buffer(False),
            ])
        image_pipeline.extend([
            ToTensor(),

• non-jit operation -> jit operation

       image_pipeline: List[Operation] = [SimpleRGBImageDecoder()]
        if name == 'train':
            image_pipeline.extend([
                Buffer(False),
                Buffer(True),
            ])
        image_pipeline.extend([
            ToTensor(),

The former case (jit operation -> non-jit operation) works correctly. But the latter case (non-jit-operation -> jit operation) failed and the first error messages were displayed.

I think that we have to fix some codes about jit compiling in loader.py to solve this issue. But I couldn't understand them by reading source codes. Do you have any ideas to fix this issue or any hints for me?

[GuillaumeLeclerc]

Hi @ar90n
Sorry I missed this. This may have been fixed in v0.0.4 ?

PS: We would really love Albumentation augmentations :)

[njwfish]

I have pulled @ar90n's fork and when I merge with master (my fork) and use albumentations in place of ffcv transforms everything seems to work!

[ar90n]

Hi @GuillaumeLeclerc
Thanks for your response. I tried the above non-jit-operation -> jit-operation code with v0.0.4. And I met the following error.

...

Exception in thread Thread-2:
Traceback (most recent call last):
  File "/opt/miniconda/envs/ffcv/lib/python3.9/threading.py", line 973, in _bootstrap_inner
    self.run()
  File "/workspaces/ffcv/ffcv/loader/epoch_iterator.py", line 83, in run
    result = self.run_pipeline(b_ix, ixes, slot, events[slot])
  File "/workspaces/ffcv/ffcv/loader/epoch_iterator.py", line 137, in run_pipeline
    result = code(*args)
  File "/opt/miniconda/envs/ffcv/lib/python3.9/site-packages/numba/core/dispatcher.py", line 468, in _compile_for_args
    error_rewrite(e, 'typing')
  File "/opt/miniconda/envs/ffcv/lib/python3.9/site-packages/numba/core/dispatcher.py", line 409, in error_rewrite
    raise e.with_traceback(None)
numba.core.errors.TypingError: Failed in nopython mode pipeline (step: nopython frontend)
non-precise type pyobject
During: typing of argument at  (2)

File "/workspaces/ffcv/examples/cifar", line 2:
<source missing, REPL/exec in use?> 

This error may have been caused by the following argument(s):
- argument 2: Cannot determine Numba type of <class 'torch.Tensor'>

After my investigation, this error was caused by the wrong synchronization between the label pipeline and the image pipeline. In the above example, the label pipeline and the image pipeline are defined as follows.

    for name in ["train", "test"]:
        label_pipeline: List[Operation] = [
            IntDecoder(),
            ToTensor(),
            ToDevice("cuda:0"),
            Squeeze(),
        ]
        image_pipeline: List[Operation] = [SimpleRGBImageDecoder()]
        if name == "train":
            image_pipeline.extend(
                [
                    Buffer(False),
                    Buffer(True),
                ]
            )
        image_pipeline.extend(
            [
                ToTensor(),
                ToDevice("cuda:0", non_blocking=True),
                ToTorchImage(),
                Convert(ch.float16),
                torchvision.transforms.Normalize(CIFAR_MEAN, CIFAR_STD),
            ]
        )

And the label pipeline is split into the following operation_blocks in the parse_pipeline method in the Pipeline Class.

0. [IntDecoder]
1. [ToTensor, ToDevice, Squeeze]

And the image pipeline is also split into the following operation_blocks.

0. [SimpleRGBImageDecoder]
1. [Buffer]
2. [Buffer]
3. [ToTensor, ToDevice, ToTorchImage, Convert, ModuleWrapper]

As you know, each operation block in the same stages is packed into the stage code in thegenerate_stage_code in the Loader class. Therefore the input of stage-2 contains torch.Tensor which is created by the stage-1 operation block of the label pipeline. And Numba can not compile torch.Tensor .

I found two workarounds. The first is the manual synchronization of pipelines. Another is making unused inputs bypass.

manual synchronization of pipelines

The motivation of this way is to synchronize the stage which converts numpy.array to torch.Tensor in label pipeline and image pipeline. This is achieved by the following. This means that two Buffer operations are inserted into the original label_pipeline.

        label_pipeline: List[Operation] = [
            IntDecoder(),
            Buffer(False),
            Buffer(True),
            ToTensor(),
            ToDevice("cuda:0"),
            Squeeze(),
        ]

As this modification, the operation_blocks of label pipeline are changed into the following.

0. [IntDecoder]
1. [Buffer]
2. [Buffer]
3. [ToTensor, ToDevice, Squeeze]

This means that conversion of numpy.array to torch.Tensor is achieved in stage-3 in label pipeline and image pipeline. So there are no torch.Tensor as the input of Numba. And this works correctly.

making unused inputs bypass

The motivation of this way is to wrap each stage codes and make unused inputs bypass original stage codes. The PoC of this workaround is following.

diff --git a/ffcv/loader/loader.py b/ffcv/loader/loader.py
index 3981da9..e044c00 100644
--- a/ffcv/loader/loader.py
+++ b/ffcv/loader/loader.py
@@ -302,7 +302,9 @@ def {fun_name}():
         function_calls = []
         memory_banks = []
         memory_banks_id = []
+        active_pipelines = set()
         for p_ix, pipeline_name, op_id, needs_indices in stage:
+            active_pipelines.add(pipeline_name)
             function_calls.append(self.generate_function_call(pipeline_name,
                                                               op_id, needs_indices))
             arg = ast.arg(arg=f'memory_{pipeline_name}_{op_id}')
@@ -316,10 +318,14 @@ def {fun_name}():

         base_code.args.args.append(ast.arg(arg='batch_indices'))

-        for p_id in self.pipelines.keys():
+        pass_through_arg_indice = []
+        arg_base_index = len(base_code.args.args)
+        for i, p_id in enumerate(self.pipelines.keys()):
             r = f'result_{p_id}'
             if stage_ix != 0:
                 base_code.args.args.append(ast.arg(arg=r))
+                if p_id not in active_pipelines:
+                    pass_through_arg_indice.append(arg_base_index + i)
             return_tuple.value.elts.append(ast.Name(id=r, ctx=ast.Load()))


@@ -341,7 +347,21 @@ def {fun_name}():
         if stage_ix % 2 == 0:
             final_code = Compiler.compile(final_code)

-        return final_code, memory_banks_id
+        def _wrap(*args):
+            args = list(args)
+            pass_through_args = []
+            for i in pass_through_arg_indice:
+                pass_through_args.append(args[i])
+                args[i] = None
+
+            ret = final_code(*args)
+            ret = list(ret)
+
+            for i, v in zip(pass_through_arg_indice, pass_through_args):
+                ret[i - arg_base_index] = v
+            return ret
+
+        return _wrap, memory_banks_id

     def generate_code(self):
         schedule = defaultdict(lambda: [])

In the above code, _wrap makes unused inputs bypass final_code. In general situation, torch.Tensor won't be input to jit-compiled stage functions. So this works correctly.

I understand this code is only PoC and this feature should be implemented as jit-compiled functions such as other stage functions.

next action

If possible, I would like to make a PR about this albumentation integration except for workarounds about torch.Tensor issue. Because this is not part of albumentation integration. What do you think?

And I want to try this torch.Tensor issue after completion of albumentation integration, if you need.

PS: I real love Albumentation augmentations, too!!

[ar90n]

@njwfish
Thanks for your report! I'm glad to hear that this modification works in another environment!

[GuillaumeLeclerc]

Seems to be resolved. Feel free to reopen if it's not the case