diff --git a/ffcv/pipeline/graph.py b/ffcv/pipeline/graph.py index ccc657bc..05da7cee 100644 --- a/ffcv/pipeline/graph.py +++ b/ffcv/pipeline/graph.py @@ -2,7 +2,12 @@ import warnings import ast -import astor +try: + # Useful for debugging + import astor +except ImportError: + pass + from collections import defaultdict from typing import Callable, Dict, List, Optional, Sequence, Set from abc import ABC, abstractmethod diff --git a/ffcv/transforms/ops.py b/ffcv/transforms/ops.py index cec29084..2b8bad5e 100644 --- a/ffcv/transforms/ops.py +++ b/ffcv/transforms/ops.py @@ -42,6 +42,8 @@ class ToDevice(Operation): def __init__(self, device, non_blocking=True): super().__init__() self.device = device + # assert isinstance(device, ch.device), \ + # f'Make sure device is a ch.device (not a {type(device)})' self.non_blocking = non_blocking def generate_code(self) -> Callable: @@ -155,4 +157,4 @@ def convert(inp, dst): return convert def declare_state_and_memory(self, previous_state: State) -> Tuple[State, Optional[AllocationQuery]]: - return replace(previous_state, dtype=self.target_dtype, jit_mode=False), None \ No newline at end of file + return replace(previous_state, dtype=self.target_dtype, jit_mode=False), None diff --git a/ffcv/transforms/translate.py b/ffcv/transforms/translate.py index e53e157e..338f47cd 100644 --- a/ffcv/transforms/translate.py +++ b/ffcv/transforms/translate.py @@ -30,18 +30,15 @@ def __init__(self, padding: int, fill: Tuple[int, int, int] = (0, 0, 0)): def generate_code(self) -> Callable: my_range = Compiler.get_iterator() pad = self.padding + fill = self.fill def translate(images, dst): n, h, w, _ = images.shape - # y_coords = randint(low=0, high=2 * pad + 1, size=(n,)) - # x_coords = randint(low=0, high=2 * pad + 1, size=(n,)) - # dst = fill - + dst[:] = fill dst[:, pad:pad+h, pad:pad+w] = images for i in my_range(n): y_coord = randint(low=0, high=2 * pad + 1) x_coord = randint(low=0, high=2 * pad + 1) - # images[i] = dst[i, y_coords[i]:y_coords[i]+h, x_coords[i]:x_coords[i]+w] images[i] = dst[i, y_coord:y_coord+h, x_coord:x_coord+w] return images