project upload

.gitattributes

@@ -34,3 +34,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text
+*.jpg filter=lfs diff=lfs merge=lfs -text
+*.jpeg filter=lfs diff=lfs merge=lfs -text

DeOldify/ImageColorizerArtisticTests.ipynb

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:48a4d84968ac49079ccddd3834d35df0b3e0e544fc687feb1d7eda4539cd5fb4
+size 508185940

DeOldify/deoldify/__init__.py

@@ -0,0 +1,8 @@
+import sys
+import logging
+logging.getLogger().addHandler(logging.StreamHandler(sys.stdout))
+logging.getLogger().setLevel(logging.INFO)
+
+from deoldify._device import _Device
+
+device = _Device()

DeOldify/deoldify/_device.py

@@ -0,0 +1,30 @@
+import os
+from enum import Enum
+from .device_id import DeviceId
+
+#NOTE:  This must be called first before any torch imports in order to work properly!
+
+class DeviceException(Exception):
+    pass
+
+class _Device:
+    def __init__(self):
+        self.set(DeviceId.CPU)
+
+    def is_gpu(self):
+        ''' Returns `True` if the current device is GPU, `False` otherwise. '''
+        return self.current() is not DeviceId.CPU
+  
+    def current(self):
+        return self._current_device
+
+    def set(self, device:DeviceId):     
+        if device == DeviceId.CPU:
+            os.environ['CUDA_VISIBLE_DEVICES']=''
+        else:
+            os.environ['CUDA_VISIBLE_DEVICES']=str(device.value)
+            import torch
+            torch.backends.cudnn.benchmark=False
+        
+        self._current_device = device    
+        return device

DeOldify/deoldify/augs.py

@@ -0,0 +1,29 @@
+import random
+
+from fastai.vision.image import TfmPixel
+
+# Contributed by Rani Horev. Thank you!
+def _noisify(
+    x, pct_pixels_min: float = 0.001, pct_pixels_max: float = 0.4, noise_range: int = 30
+):
+    if noise_range > 255 or noise_range < 0:
+        raise Exception("noise_range must be between 0 and 255, inclusively.")
+
+    h, w = x.shape[1:]
+    img_size = h * w
+    mult = 10000.0
+    pct_pixels = (
+        random.randrange(int(pct_pixels_min * mult), int(pct_pixels_max * mult)) / mult
+    )
+    noise_count = int(img_size * pct_pixels)
+
+    for ii in range(noise_count):
+        yy = random.randrange(h)
+        xx = random.randrange(w)
+        noise = random.randrange(-noise_range, noise_range) / 255.0
+        x[:, yy, xx].add_(noise)
+
+    return x
+
+
+noisify = TfmPixel(_noisify)

DeOldify/deoldify/critics.py

@@ -0,0 +1,47 @@
+from fastai.basic_train import Learner
+from fastai.core import *
+from fastai.layers import NormType, conv_layer
+from fastai.torch_core import *
+from fastai.vision import *
+from fastai.vision.data import ImageDataBunch
+from fastai.vision.gan import AdaptiveLoss, accuracy_thresh_expand
+
+_conv_args = dict(leaky=0.2, norm_type=NormType.Spectral)
+
+
+def _conv(ni: int, nf: int, ks: int = 3, stride: int = 1, **kwargs):
+    return conv_layer(ni, nf, ks=ks, stride=stride, **_conv_args, **kwargs)
+
+
+def custom_gan_critic(
+    n_channels: int = 3, nf: int = 256, n_blocks: int = 3, p: int = 0.15
+):
+    "Critic to train a `GAN`."
+    layers = [_conv(n_channels, nf, ks=4, stride=2), nn.Dropout2d(p / 2)]
+    for i in range(n_blocks):
+        layers += [
+            _conv(nf, nf, ks=3, stride=1),
+            nn.Dropout2d(p),
+            _conv(nf, nf * 2, ks=4, stride=2, self_attention=(i == 0)),
+        ]
+        nf *= 2
+    layers += [
+        _conv(nf, nf, ks=3, stride=1),
+        _conv(nf, 1, ks=4, bias=False, padding=0, use_activ=False),
+        Flatten(),
+    ]
+    return nn.Sequential(*layers)
+
+
+def colorize_crit_learner(
+    data: ImageDataBunch,
+    loss_critic=AdaptiveLoss(nn.BCEWithLogitsLoss()),
+    nf: int = 256,
+) -> Learner:
+    return Learner(
+        data,
+        custom_gan_critic(nf=nf),
+        metrics=accuracy_thresh_expand,
+        loss_func=loss_critic,
+        wd=1e-3,
+    )

DeOldify/deoldify/dataset.py

@@ -0,0 +1,46 @@
+from fastai import *
+from fastai.core import *
+from fastai.vision.transform import get_transforms
+from fastai.vision.data import ImageImageList, ImageDataBunch, imagenet_stats
+
+
+def get_colorize_data(
+    sz: int,
+    bs: int,
+    crappy_path: Path,
+    good_path: Path,
+    random_seed: int = None,
+    keep_pct: float = 1.0,
+    num_workers: int = 8,
+    stats: tuple = imagenet_stats,
+    xtra_tfms=[],
+) -> ImageDataBunch:
+    
+    src = (
+        ImageImageList.from_folder(crappy_path, convert_mode='RGB')
+        .use_partial_data(sample_pct=keep_pct, seed=random_seed)
+        .split_by_rand_pct(0.1, seed=random_seed)
+    )
+
+    data = (
+        src.label_from_func(lambda x: good_path / x.relative_to(crappy_path))
+        .transform(
+            get_transforms(
+                max_zoom=1.2, max_lighting=0.5, max_warp=0.25, xtra_tfms=xtra_tfms
+            ),
+            size=sz,
+            tfm_y=True,
+        )
+        .databunch(bs=bs, num_workers=num_workers, no_check=True)
+        .normalize(stats, do_y=True)
+    )
+
+    data.c = 3
+    return data
+
+
+def get_dummy_databunch() -> ImageDataBunch:
+    path = Path('./dummy/')
+    return get_colorize_data(
+        sz=1, bs=1, crappy_path=path, good_path=path, keep_pct=0.001
+    )

DeOldify/deoldify/device_id.py

@@ -0,0 +1,12 @@
+from enum import IntEnum
+
+class DeviceId(IntEnum):
+    GPU0 = 0,
+    GPU1 = 1,
+    GPU2 = 2,
+    GPU3 = 3,
+    GPU4 = 4,
+    GPU5 = 5,
+    GPU6 = 6,
+    GPU7 = 7,
+    CPU = 99

DeOldify/deoldify/filters.py

@@ -0,0 +1,124 @@
+from fastai.basic_data import DatasetType
+from fastai.basic_train import Learner
+from abc import ABC, abstractmethod
+from fastai.core import *
+from fastai.vision import *
+from fastai.vision.image import *
+from fastai.vision.data import *
+from fastai import *
+import cv2
+from PIL import Image as PilImage
+from deoldify import device as device_settings
+import logging
+
+
+class IFilter(ABC):
+    @abstractmethod
+    def filter(
+        self, orig_image: PilImage, filtered_image: PilImage, render_factor: int
+    ) -> PilImage:
+        pass
+
+
+class BaseFilter(IFilter):
+    def __init__(self, learn: Learner, stats: tuple = imagenet_stats):
+        super().__init__()
+        self.learn = learn
+        
+        if not device_settings.is_gpu():
+            self.learn.model = self.learn.model.cpu()
+        
+        self.device = next(self.learn.model.parameters()).device
+        self.norm, self.denorm = normalize_funcs(*stats)
+
+    def _transform(self, image: PilImage) -> PilImage:
+        return image
+
+    def _scale_to_square(self, orig: PilImage, targ: int) -> PilImage:
+        # a simple stretch to fit a square really makes a big difference in rendering quality/consistency.
+        # I've tried padding to the square as well (reflect, symetric, constant, etc).  Not as good!
+        targ_sz = (targ, targ)
+        return orig.resize(targ_sz, resample=PIL.Image.BILINEAR)
+
+    def _get_model_ready_image(self, orig: PilImage, sz: int) -> PilImage:
+        result = self._scale_to_square(orig, sz)
+        result = self._transform(result)
+        return result
+
+    def _model_process(self, orig: PilImage, sz: int) -> PilImage:
+        model_image = self._get_model_ready_image(orig, sz)
+        x = pil2tensor(model_image, np.float32)
+        x = x.to(self.device)
+        x.div_(255)
+        x, y = self.norm((x, x), do_x=True)
+        
+        try:
+            result = self.learn.pred_batch(
+                ds_type=DatasetType.Valid, batch=(x[None], y[None]), reconstruct=True
+            )
+        except RuntimeError as rerr:
+            if 'memory' not in str(rerr):
+                raise rerr
+            logging.warn('Warning: render_factor was set too high, and out of memory error resulted. Returning original image.')
+            return model_image
+            
+        out = result[0]
+        out = self.denorm(out.px, do_x=False)
+        out = image2np(out * 255).astype(np.uint8)
+        return PilImage.fromarray(out)
+
+    def _unsquare(self, image: PilImage, orig: PilImage) -> PilImage:
+        targ_sz = orig.size
+        image = image.resize(targ_sz, resample=PIL.Image.BILINEAR)
+        return image
+
+
+class ColorizerFilter(BaseFilter):
+    def __init__(self, learn: Learner, stats: tuple = imagenet_stats):
+        super().__init__(learn=learn, stats=stats)
+        self.render_base = 16
+
+    def filter(
+        self, orig_image: PilImage, filtered_image: PilImage, render_factor: int, post_process: bool = True) -> PilImage:
+        render_sz = render_factor * self.render_base
+        model_image = self._model_process(orig=filtered_image, sz=render_sz)
+        raw_color = self._unsquare(model_image, orig_image)
+
+        if post_process:
+            return self._post_process(raw_color, orig_image)
+        else:
+            return raw_color
+
+    def _transform(self, image: PilImage) -> PilImage:
+        return image.convert('LA').convert('RGB')
+
+    # This takes advantage of the fact that human eyes are much less sensitive to
+    # imperfections in chrominance compared to luminance.  This means we can
+    # save a lot on memory and processing in the model, yet get a great high
+    # resolution result at the end.  This is primarily intended just for
+    # inference
+    def _post_process(self, raw_color: PilImage, orig: PilImage) -> PilImage:
+        color_np = np.asarray(raw_color)
+        orig_np = np.asarray(orig)
+        color_yuv = cv2.cvtColor(color_np, cv2.COLOR_RGB2YUV)
+        # do a black and white transform first to get better luminance values
+        orig_yuv = cv2.cvtColor(orig_np, cv2.COLOR_RGB2YUV)
+        hires = np.copy(orig_yuv)
+        hires[:, :, 1:3] = color_yuv[:, :, 1:3]
+        final = cv2.cvtColor(hires, cv2.COLOR_YUV2RGB)
+        final = PilImage.fromarray(final)
+        return final
+
+
+class MasterFilter(BaseFilter):
+    def __init__(self, filters: List[IFilter], render_factor: int):
+        self.filters = filters
+        self.render_factor = render_factor
+
+    def filter(
+        self, orig_image: PilImage, filtered_image: PilImage, render_factor: int = None, post_process: bool = True) -> PilImage:
+        render_factor = self.render_factor if render_factor is None else render_factor
+        for filter in self.filters:
+            filtered_image = filter.filter(orig_image, filtered_image, render_factor, post_process)
+
+        return filtered_image

DeOldify/deoldify/generators.py

@@ -0,0 +1,155 @@
+from fastai.basic_data import DataBunch
+from fastai.basic_train import Learner
+from fastai.layers import NormType
+from fastai.torch_core import SplitFuncOrIdxList, apply_init, to_device
+from fastai.vision import *
+from fastai.vision.learner import cnn_config, create_body
+from torch import nn
+from .unet import DynamicUnetWide, DynamicUnetDeep
+from .dataset import *
+
+# Weights are implicitly read from ./models/ folder
+def gen_inference_wide(
+    root_folder: Path, weights_name: str, nf_factor: int = 2, arch=models.resnet101) -> Learner:
+    data = get_dummy_databunch()
+    learn = gen_learner_wide(
+        data=data, gen_loss=F.l1_loss, nf_factor=nf_factor, arch=arch
+    )
+    learn.path = root_folder
+    learn.load(weights_name)
+    learn.model.eval()
+    return learn
+
+
+def gen_learner_wide(
+    data: ImageDataBunch, gen_loss, arch=models.resnet101, nf_factor: int = 2
+) -> Learner:
+    return unet_learner_wide(
+        data,
+        arch=arch,
+        wd=1e-3,
+        blur=True,
+        norm_type=NormType.Spectral,
+        self_attention=True,
+        y_range=(-3.0, 3.0),
+        loss_func=gen_loss,
+        nf_factor=nf_factor,
+    )
+
+
+# The code below is meant to be merged into fastaiv1 ideally
+def unet_learner_wide(
+    data: DataBunch,
+    arch: Callable,
+    pretrained: bool = True,
+    blur_final: bool = True,
+    norm_type: Optional[NormType] = NormType,
+    split_on: Optional[SplitFuncOrIdxList] = None,
+    blur: bool = False,
+    self_attention: bool = False,
+    y_range: Optional[Tuple[float, float]] = None,
+    last_cross: bool = True,
+    bottle: bool = False,
+    nf_factor: int = 1,
+    **kwargs: Any
+) -> Learner:
+    "Build Unet learner from `data` and `arch`."
+    meta = cnn_config(arch)
+    body = create_body(arch, pretrained)
+    model = to_device(
+        DynamicUnetWide(
+            body,
+            n_classes=data.c,
+            blur=blur,
+            blur_final=blur_final,
+            self_attention=self_attention,
+            y_range=y_range,
+            norm_type=norm_type,
+            last_cross=last_cross,
+            bottle=bottle,
+            nf_factor=nf_factor,
+        ),
+        data.device,
+    )
+    learn = Learner(data, model, **kwargs)
+    learn.split(ifnone(split_on, meta['split']))
+    if pretrained:
+        learn.freeze()
+    apply_init(model[2], nn.init.kaiming_normal_)
+    return learn
+
+
+# ----------------------------------------------------------------------
+
+# Weights are implicitly read from ./models/ folder
+def gen_inference_deep(
+    root_folder: Path, weights_name: str, arch=models.resnet34, nf_factor: float = 1.5) -> Learner:
+    data = get_dummy_databunch()
+    learn = gen_learner_deep(
+        data=data, gen_loss=F.l1_loss, arch=arch, nf_factor=nf_factor
+    )
+    learn.path = root_folder
+    learn.load(weights_name)
+    learn.model.eval()
+    return learn
+
+
+def gen_learner_deep(
+    data: ImageDataBunch, gen_loss, arch=models.resnet34, nf_factor: float = 1.5
+) -> Learner:
+    return unet_learner_deep(
+        data,
+        arch,
+        wd=1e-3,
+        blur=True,
+        norm_type=NormType.Spectral,
+        self_attention=True,
+        y_range=(-3.0, 3.0),
+        loss_func=gen_loss,
+        nf_factor=nf_factor,
+    )
+
+
+# The code below is meant to be merged into fastaiv1 ideally
+def unet_learner_deep(
+    data: DataBunch,
+    arch: Callable,
+    pretrained: bool = True,
+    blur_final: bool = True,
+    norm_type: Optional[NormType] = NormType,
+    split_on: Optional[SplitFuncOrIdxList] = None,
+    blur: bool = False,
+    self_attention: bool = False,
+    y_range: Optional[Tuple[float, float]] = None,
+    last_cross: bool = True,
+    bottle: bool = False,
+    nf_factor: float = 1.5,
+    **kwargs: Any
+) -> Learner:
+    "Build Unet learner from `data` and `arch`."
+    meta = cnn_config(arch)
+    body = create_body(arch, pretrained)
+    model = to_device(
+        DynamicUnetDeep(
+            body,
+            n_classes=data.c,
+            blur=blur,
+            blur_final=blur_final,
+            self_attention=self_attention,
+            y_range=y_range,
+            norm_type=norm_type,
+            last_cross=last_cross,
+            bottle=bottle,
+            nf_factor=nf_factor,
+        ),
+        data.device,
+    )
+    learn = Learner(data, model, **kwargs)
+    learn.split(ifnone(split_on, meta['split']))
+    if pretrained:
+        learn.freeze()
+    apply_init(model[2], nn.init.kaiming_normal_)
+    return learn
+
+
+# -----------------------------

DeOldify/deoldify/layers.py

@@ -0,0 +1,46 @@
+from fastai.layers import *
+from fastai.torch_core import *
+
+
+# The code below is meant to be merged into fastaiv1 ideally
+
+
+def custom_conv_layer(
+    ni: int,
+    nf: int,
+    ks: int = 3,
+    stride: int = 1,
+    padding: int = None,
+    bias: bool = None,
+    is_1d: bool = False,
+    norm_type: Optional[NormType] = NormType.Batch,
+    use_activ: bool = True,
+    leaky: float = None,
+    transpose: bool = False,
+    init: Callable = nn.init.kaiming_normal_,
+    self_attention: bool = False,
+    extra_bn: bool = False,
+):
+    "Create a sequence of convolutional (`ni` to `nf`), ReLU (if `use_activ`) and batchnorm (if `bn`) layers."
+    if padding is None:
+        padding = (ks - 1) // 2 if not transpose else 0
+    bn = norm_type in (NormType.Batch, NormType.BatchZero) or extra_bn == True
+    if bias is None:
+        bias = not bn
+    conv_func = nn.ConvTranspose2d if transpose else nn.Conv1d if is_1d else nn.Conv2d
+    conv = init_default(
+        conv_func(ni, nf, kernel_size=ks, bias=bias, stride=stride, padding=padding),
+        init,
+    )
+    if norm_type == NormType.Weight:
+        conv = weight_norm(conv)
+    elif norm_type == NormType.Spectral:
+        conv = spectral_norm(conv)
+    layers = [conv]
+    if use_activ:
+        layers.append(relu(True, leaky=leaky))
+    if bn:
+        layers.append((nn.BatchNorm1d if is_1d else nn.BatchNorm2d)(nf))
+    if self_attention:
+        layers.append(SelfAttention(nf))
+    return nn.Sequential(*layers)

DeOldify/deoldify/loss.py

@@ -0,0 +1,136 @@
+from fastai import *
+from fastai.core import *
+from fastai.torch_core import *
+from fastai.callbacks import hook_outputs
+import torchvision.models as models
+
+
+class FeatureLoss(nn.Module):
+    def __init__(self, layer_wgts=[20, 70, 10]):
+        super().__init__()
+
+        self.m_feat = models.vgg16_bn(True).features.cuda().eval()
+        requires_grad(self.m_feat, False)
+        blocks = [
+            i - 1
+            for i, o in enumerate(children(self.m_feat))
+            if isinstance(o, nn.MaxPool2d)
+        ]
+        layer_ids = blocks[2:5]
+        self.loss_features = [self.m_feat[i] for i in layer_ids]
+        self.hooks = hook_outputs(self.loss_features, detach=False)
+        self.wgts = layer_wgts
+        self.metric_names = ['pixel'] + [f'feat_{i}' for i in range(len(layer_ids))]
+        self.base_loss = F.l1_loss
+
+    def _make_features(self, x, clone=False):
+        self.m_feat(x)
+        return [(o.clone() if clone else o) for o in self.hooks.stored]
+
+    def forward(self, input, target):
+        out_feat = self._make_features(target, clone=True)
+        in_feat = self._make_features(input)
+        self.feat_losses = [self.base_loss(input, target)]
+        self.feat_losses += [
+            self.base_loss(f_in, f_out) * w
+            for f_in, f_out, w in zip(in_feat, out_feat, self.wgts)
+        ]
+
+        self.metrics = dict(zip(self.metric_names, self.feat_losses))
+        return sum(self.feat_losses)
+
+    def __del__(self):
+        self.hooks.remove()
+
+
+# Refactored code, originally from https://github.com/VinceMarron/style_transfer
+class WassFeatureLoss(nn.Module):
+    def __init__(self, layer_wgts=[5, 15, 2], wass_wgts=[3.0, 0.7, 0.01]):
+        super().__init__()
+        self.m_feat = models.vgg16_bn(True).features.cuda().eval()
+        requires_grad(self.m_feat, False)
+        blocks = [
+            i - 1
+            for i, o in enumerate(children(self.m_feat))
+            if isinstance(o, nn.MaxPool2d)
+        ]
+        layer_ids = blocks[2:5]
+        self.loss_features = [self.m_feat[i] for i in layer_ids]
+        self.hooks = hook_outputs(self.loss_features, detach=False)
+        self.wgts = layer_wgts
+        self.wass_wgts = wass_wgts
+        self.metric_names = (
+            ['pixel']
+            + [f'feat_{i}' for i in range(len(layer_ids))]
+            + [f'wass_{i}' for i in range(len(layer_ids))]
+        )
+        self.base_loss = F.l1_loss
+
+    def _make_features(self, x, clone=False):
+        self.m_feat(x)
+        return [(o.clone() if clone else o) for o in self.hooks.stored]
+
+    def _calc_2_moments(self, tensor):
+        chans = tensor.shape[1]
+        tensor = tensor.view(1, chans, -1)
+        n = tensor.shape[2]
+        mu = tensor.mean(2)
+        tensor = (tensor - mu[:, :, None]).squeeze(0)
+        # Prevents nasty bug that happens very occassionally- divide by zero.  Why such things happen?
+        if n == 0:
+            return None, None
+        cov = torch.mm(tensor, tensor.t()) / float(n)
+        return mu, cov
+
+    def _get_style_vals(self, tensor):
+        mean, cov = self._calc_2_moments(tensor)
+        if mean is None:
+            return None, None, None
+        eigvals, eigvects = torch.symeig(cov, eigenvectors=True)
+        eigroot_mat = torch.diag(torch.sqrt(eigvals.clamp(min=0)))
+        root_cov = torch.mm(torch.mm(eigvects, eigroot_mat), eigvects.t())
+        tr_cov = eigvals.clamp(min=0).sum()
+        return mean, tr_cov, root_cov
+
+    def _calc_l2wass_dist(
+        self, mean_stl, tr_cov_stl, root_cov_stl, mean_synth, cov_synth
+    ):
+        tr_cov_synth = torch.symeig(cov_synth, eigenvectors=True)[0].clamp(min=0).sum()
+        mean_diff_squared = (mean_stl - mean_synth).pow(2).sum()
+        cov_prod = torch.mm(torch.mm(root_cov_stl, cov_synth), root_cov_stl)
+        var_overlap = torch.sqrt(
+            torch.symeig(cov_prod, eigenvectors=True)[0].clamp(min=0) + 1e-8
+        ).sum()
+        dist = mean_diff_squared + tr_cov_stl + tr_cov_synth - 2 * var_overlap
+        return dist
+
+    def _single_wass_loss(self, pred, targ):
+        mean_test, tr_cov_test, root_cov_test = targ
+        mean_synth, cov_synth = self._calc_2_moments(pred)
+        loss = self._calc_l2wass_dist(
+            mean_test, tr_cov_test, root_cov_test, mean_synth, cov_synth
+        )
+        return loss
+
+    def forward(self, input, target):
+        out_feat = self._make_features(target, clone=True)
+        in_feat = self._make_features(input)
+        self.feat_losses = [self.base_loss(input, target)]
+        self.feat_losses += [
+            self.base_loss(f_in, f_out) * w
+            for f_in, f_out, w in zip(in_feat, out_feat, self.wgts)
+        ]
+
+        styles = [self._get_style_vals(i) for i in out_feat]
+
+        if styles[0][0] is not None:
+            self.feat_losses += [
+                self._single_wass_loss(f_pred, f_targ) * w
+                for f_pred, f_targ, w in zip(in_feat, styles, self.wass_wgts)
+            ]
+
+        self.metrics = dict(zip(self.metric_names, self.feat_losses))
+        return sum(self.feat_losses)
+
+    def __del__(self):
+        self.hooks.remove()

DeOldify/deoldify/save.py

@@ -0,0 +1,29 @@
+from fastai.basic_train import Learner, LearnerCallback
+from fastai.vision.gan import GANLearner
+
+
+class GANSaveCallback(LearnerCallback):
+    """A `LearnerCallback` that saves history of metrics while training `learn` into CSV `filename`."""
+
+    def __init__(
+        self,
+        learn: GANLearner,
+        learn_gen: Learner,
+        filename: str,
+        save_iters: int = 1000,
+    ):
+        super().__init__(learn)
+        self.learn_gen = learn_gen
+        self.filename = filename
+        self.save_iters = save_iters
+
+    def on_batch_end(self, iteration: int, epoch: int, **kwargs) -> None:
+        if iteration == 0:
+            return
+
+        if iteration % self.save_iters == 0:
+            self._save_gen_learner(iteration=iteration, epoch=epoch)
+
+    def _save_gen_learner(self, iteration: int, epoch: int):
+        filename = '{}_{}_{}'.format(self.filename, epoch, iteration)
+        self.learn_gen.save(filename)

DeOldify/deoldify/unet.py

@@ -0,0 +1,285 @@
+from fastai.layers import *
+from .layers import *
+from fastai.torch_core import *
+from fastai.callbacks.hooks import *
+from fastai.vision import *
+
+
+# The code below is meant to be merged into fastaiv1 ideally
+
+__all__ = ['DynamicUnetDeep', 'DynamicUnetWide']
+
+
+def _get_sfs_idxs(sizes: Sizes) -> List[int]:
+    "Get the indexes of the layers where the size of the activation changes."
+    feature_szs = [size[-1] for size in sizes]
+    sfs_idxs = list(
+        np.where(np.array(feature_szs[:-1]) != np.array(feature_szs[1:]))[0]
+    )
+    if feature_szs[0] != feature_szs[1]:
+        sfs_idxs = [0] + sfs_idxs
+    return sfs_idxs
+
+
+class CustomPixelShuffle_ICNR(nn.Module):
+    "Upsample by `scale` from `ni` filters to `nf` (default `ni`), using `nn.PixelShuffle`, `icnr` init, and `weight_norm`."
+
+    def __init__(
+        self,
+        ni: int,
+        nf: int = None,
+        scale: int = 2,
+        blur: bool = False,
+        leaky: float = None,
+        **kwargs
+    ):
+        super().__init__()
+        nf = ifnone(nf, ni)
+        self.conv = custom_conv_layer(
+            ni, nf * (scale ** 2), ks=1, use_activ=False, **kwargs
+        )
+        icnr(self.conv[0].weight)
+        self.shuf = nn.PixelShuffle(scale)
+        # Blurring over (h*w) kernel
+        # "Super-Resolution using Convolutional Neural Networks without Any Checkerboard Artifacts"
+        # - https://arxiv.org/abs/1806.02658
+        self.pad = nn.ReplicationPad2d((1, 0, 1, 0))
+        self.blur = nn.AvgPool2d(2, stride=1)
+        self.relu = relu(True, leaky=leaky)
+
+    def forward(self, x):
+        x = self.shuf(self.relu(self.conv(x)))
+        return self.blur(self.pad(x)) if self.blur else x
+
+
+class UnetBlockDeep(nn.Module):
+    "A quasi-UNet block, using `PixelShuffle_ICNR upsampling`."
+
+    def __init__(
+        self,
+        up_in_c: int,
+        x_in_c: int,
+        hook: Hook,
+        final_div: bool = True,
+        blur: bool = False,
+        leaky: float = None,
+        self_attention: bool = False,
+        nf_factor: float = 1.0,
+        **kwargs
+    ):
+        super().__init__()
+        self.hook = hook
+        self.shuf = CustomPixelShuffle_ICNR(
+            up_in_c, up_in_c // 2, blur=blur, leaky=leaky, **kwargs
+        )
+        self.bn = batchnorm_2d(x_in_c)
+        ni = up_in_c // 2 + x_in_c
+        nf = int((ni if final_div else ni // 2) * nf_factor)
+        self.conv1 = custom_conv_layer(ni, nf, leaky=leaky, **kwargs)
+        self.conv2 = custom_conv_layer(
+            nf, nf, leaky=leaky, self_attention=self_attention, **kwargs
+        )
+        self.relu = relu(leaky=leaky)
+
+    def forward(self, up_in: Tensor) -> Tensor:
+        s = self.hook.stored
+        up_out = self.shuf(up_in)
+        ssh = s.shape[-2:]
+        if ssh != up_out.shape[-2:]:
+            up_out = F.interpolate(up_out, s.shape[-2:], mode='nearest')
+        cat_x = self.relu(torch.cat([up_out, self.bn(s)], dim=1))
+        return self.conv2(self.conv1(cat_x))
+
+
+class DynamicUnetDeep(SequentialEx):
+    "Create a U-Net from a given architecture."
+
+    def __init__(
+        self,
+        encoder: nn.Module,
+        n_classes: int,
+        blur: bool = False,
+        blur_final=True,
+        self_attention: bool = False,
+        y_range: Optional[Tuple[float, float]] = None,
+        last_cross: bool = True,
+        bottle: bool = False,
+        norm_type: Optional[NormType] = NormType.Batch,
+        nf_factor: float = 1.0,
+        **kwargs
+    ):
+        extra_bn = norm_type == NormType.Spectral
+        imsize = (256, 256)
+        sfs_szs = model_sizes(encoder, size=imsize)
+        sfs_idxs = list(reversed(_get_sfs_idxs(sfs_szs)))
+        self.sfs = hook_outputs([encoder[i] for i in sfs_idxs], detach=False)
+        x = dummy_eval(encoder, imsize).detach()
+
+        ni = sfs_szs[-1][1]
+        middle_conv = nn.Sequential(
+            custom_conv_layer(
+                ni, ni * 2, norm_type=norm_type, extra_bn=extra_bn, **kwargs
+            ),
+            custom_conv_layer(
+                ni * 2, ni, norm_type=norm_type, extra_bn=extra_bn, **kwargs
+            ),
+        ).eval()
+        x = middle_conv(x)
+        layers = [encoder, batchnorm_2d(ni), nn.ReLU(), middle_conv]
+
+        for i, idx in enumerate(sfs_idxs):
+            not_final = i != len(sfs_idxs) - 1
+            up_in_c, x_in_c = int(x.shape[1]), int(sfs_szs[idx][1])
+            do_blur = blur and (not_final or blur_final)
+            sa = self_attention and (i == len(sfs_idxs) - 3)
+            unet_block = UnetBlockDeep(
+                up_in_c,
+                x_in_c,
+                self.sfs[i],
+                final_div=not_final,
+                blur=blur,
+                self_attention=sa,
+                norm_type=norm_type,
+                extra_bn=extra_bn,
+                nf_factor=nf_factor,
+                **kwargs
+            ).eval()
+            layers.append(unet_block)
+            x = unet_block(x)
+
+        ni = x.shape[1]
+        if imsize != sfs_szs[0][-2:]:
+            layers.append(PixelShuffle_ICNR(ni, **kwargs))
+        if last_cross:
+            layers.append(MergeLayer(dense=True))
+            ni += in_channels(encoder)
+            layers.append(res_block(ni, bottle=bottle, norm_type=norm_type, **kwargs))
+        layers += [
+            custom_conv_layer(ni, n_classes, ks=1, use_activ=False, norm_type=norm_type)
+        ]
+        if y_range is not None:
+            layers.append(SigmoidRange(*y_range))
+        super().__init__(*layers)
+
+    def __del__(self):
+        if hasattr(self, "sfs"):
+            self.sfs.remove()
+
+
+# ------------------------------------------------------
+class UnetBlockWide(nn.Module):
+    "A quasi-UNet block, using `PixelShuffle_ICNR upsampling`."
+
+    def __init__(
+        self,
+        up_in_c: int,
+        x_in_c: int,
+        n_out: int,
+        hook: Hook,
+        final_div: bool = True,
+        blur: bool = False,
+        leaky: float = None,
+        self_attention: bool = False,
+        **kwargs
+    ):
+        super().__init__()
+        self.hook = hook
+        up_out = x_out = n_out // 2
+        self.shuf = CustomPixelShuffle_ICNR(
+            up_in_c, up_out, blur=blur, leaky=leaky, **kwargs
+        )
+        self.bn = batchnorm_2d(x_in_c)
+        ni = up_out + x_in_c
+        self.conv = custom_conv_layer(
+            ni, x_out, leaky=leaky, self_attention=self_attention, **kwargs
+        )
+        self.relu = relu(leaky=leaky)
+
+    def forward(self, up_in: Tensor) -> Tensor:
+        s = self.hook.stored
+        up_out = self.shuf(up_in)
+        ssh = s.shape[-2:]
+        if ssh != up_out.shape[-2:]:
+            up_out = F.interpolate(up_out, s.shape[-2:], mode='nearest')
+        cat_x = self.relu(torch.cat([up_out, self.bn(s)], dim=1))
+        return self.conv(cat_x)
+
+
+class DynamicUnetWide(SequentialEx):
+    "Create a U-Net from a given architecture."
+
+    def __init__(
+        self,
+        encoder: nn.Module,
+        n_classes: int,
+        blur: bool = False,
+        blur_final=True,
+        self_attention: bool = False,
+        y_range: Optional[Tuple[float, float]] = None,
+        last_cross: bool = True,
+        bottle: bool = False,
+        norm_type: Optional[NormType] = NormType.Batch,
+        nf_factor: int = 1,
+        **kwargs
+    ):
+
+        nf = 512 * nf_factor
+        extra_bn = norm_type == NormType.Spectral
+        imsize = (256, 256)
+        sfs_szs = model_sizes(encoder, size=imsize)
+        sfs_idxs = list(reversed(_get_sfs_idxs(sfs_szs)))
+        self.sfs = hook_outputs([encoder[i] for i in sfs_idxs], detach=False)
+        x = dummy_eval(encoder, imsize).detach()
+
+        ni = sfs_szs[-1][1]
+        middle_conv = nn.Sequential(
+            custom_conv_layer(
+                ni, ni * 2, norm_type=norm_type, extra_bn=extra_bn, **kwargs
+            ),
+            custom_conv_layer(
+                ni * 2, ni, norm_type=norm_type, extra_bn=extra_bn, **kwargs
+            ),
+        ).eval()
+        x = middle_conv(x)
+        layers = [encoder, batchnorm_2d(ni), nn.ReLU(), middle_conv]
+
+        for i, idx in enumerate(sfs_idxs):
+            not_final = i != len(sfs_idxs) - 1
+            up_in_c, x_in_c = int(x.shape[1]), int(sfs_szs[idx][1])
+            do_blur = blur and (not_final or blur_final)
+            sa = self_attention and (i == len(sfs_idxs) - 3)
+
+            n_out = nf if not_final else nf // 2
+
+            unet_block = UnetBlockWide(
+                up_in_c,
+                x_in_c,
+                n_out,
+                self.sfs[i],
+                final_div=not_final,
+                blur=blur,
+                self_attention=sa,
+                norm_type=norm_type,
+                extra_bn=extra_bn,
+                **kwargs
+            ).eval()
+            layers.append(unet_block)
+            x = unet_block(x)
+
+        ni = x.shape[1]
+        if imsize != sfs_szs[0][-2:]:
+            layers.append(PixelShuffle_ICNR(ni, **kwargs))
+        if last_cross:
+            layers.append(MergeLayer(dense=True))
+            ni += in_channels(encoder)
+            layers.append(res_block(ni, bottle=bottle, norm_type=norm_type, **kwargs))
+        layers += [
+            custom_conv_layer(ni, n_classes, ks=1, use_activ=False, norm_type=norm_type)
+        ]
+        if y_range is not None:
+            layers.append(SigmoidRange(*y_range))
+        super().__init__(*layers)
+
+    def __del__(self):
+        if hasattr(self, "sfs"):
+            self.sfs.remove()

DeOldify/deoldify/visualize.py

@@ -0,0 +1,487 @@
+from fastai.core import *
+from fastai.vision import *
+from matplotlib.axes import Axes
+from .filters import IFilter, MasterFilter, ColorizerFilter
+from .generators import gen_inference_deep, gen_inference_wide
+from PIL import Image
+import ffmpeg
+import yt_dlp as youtube_dl
+import gc
+import requests
+from io import BytesIO
+import base64
+from IPython import display as ipythondisplay
+from IPython.display import HTML
+from IPython.display import Image as ipythonimage
+import cv2
+import logging
+
+# adapted from https://www.pyimagesearch.com/2016/04/25/watermarking-images-with-opencv-and-python/
+def get_watermarked(pil_image: Image) -> Image:
+    try:
+        image = cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)
+        (h, w) = image.shape[:2]
+        image = np.dstack([image, np.ones((h, w), dtype="uint8") * 255])
+        pct = 0.05
+        full_watermark = cv2.imread(
+            './resource_images/watermark.png', cv2.IMREAD_UNCHANGED
+        )
+        (fwH, fwW) = full_watermark.shape[:2]
+        wH = int(pct * h)
+        wW = int((pct * h / fwH) * fwW)
+        watermark = cv2.resize(full_watermark, (wH, wW), interpolation=cv2.INTER_AREA)
+        overlay = np.zeros((h, w, 4), dtype="uint8")
+        (wH, wW) = watermark.shape[:2]
+        overlay[h - wH - 10 : h - 10, 10 : 10 + wW] = watermark
+        # blend the two images together using transparent overlays
+        output = image.copy()
+        cv2.addWeighted(overlay, 0.5, output, 1.0, 0, output)
+        rgb_image = cv2.cvtColor(output, cv2.COLOR_BGR2RGB)
+        final_image = Image.fromarray(rgb_image)
+        return final_image
+    except:
+        # Don't want this to crash everything, so let's just not watermark the image for now.
+        return pil_image
+
+
+class ModelImageVisualizer:
+    def __init__(self, filter: IFilter, results_dir: str = None):
+        self.filter = filter
+        self.results_dir = None if results_dir is None else Path(results_dir)
+        self.results_dir.mkdir(parents=True, exist_ok=True)
+
+    def _clean_mem(self):
+        torch.cuda.empty_cache()
+        # gc.collect()
+
+    def _open_pil_image(self, path: Path) -> Image:
+        return PIL.Image.open(path).convert('RGB')
+
+    def _get_image_from_url(self, url: str) -> Image:
+        response = requests.get(url, timeout=30, headers={'user-agent':'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/62.0.3202.94 Safari/537.36'})
+        img = PIL.Image.open(BytesIO(response.content)).convert('RGB')
+        return img
+
+    def plot_transformed_image_from_url(
+        self,
+        url: str,
+        path: str = 'test_images/image.png',
+        results_dir:Path = None,
+        figsize: Tuple[int, int] = (20, 20),
+        render_factor: int = None,
+        
+        display_render_factor: bool = False,
+        compare: bool = False,
+        post_process: bool = True,
+        watermarked: bool = True,
+    ) -> Path:
+        img = self._get_image_from_url(url)
+        img.save(path)
+        return self.plot_transformed_image(
+            path=path,
+            results_dir=results_dir,
+            figsize=figsize,
+            render_factor=render_factor,
+            display_render_factor=display_render_factor,
+            compare=compare,
+            post_process = post_process,
+            watermarked=watermarked,
+        )
+
+    def plot_transformed_image(
+        self,
+        path: str,
+        results_dir:Path = None,
+        figsize: Tuple[int, int] = (20, 20),
+        render_factor: int = None,
+        display_render_factor: bool = False,
+        compare: bool = False,
+        post_process: bool = True,
+        watermarked: bool = True,
+    ) -> Path:
+        path = Path(path)
+        if results_dir is None:
+            results_dir = Path(self.results_dir)
+        result = self.get_transformed_image(
+            path, render_factor, post_process=post_process,watermarked=watermarked
+        )
+        orig = self._open_pil_image(path)
+        if compare:
+            self._plot_comparison(
+                figsize, render_factor, display_render_factor, orig, result
+            )
+        else:
+            self._plot_solo(figsize, render_factor, display_render_factor, result)
+
+        orig.close()
+        result_path = self._save_result_image(path, result, results_dir=results_dir)
+        result.close()
+        return result_path
+
+    def _plot_comparison(
+        self,
+        figsize: Tuple[int, int],
+        render_factor: int,
+        display_render_factor: bool,
+        orig: Image,
+        result: Image,
+    ):
+        fig, axes = plt.subplots(1, 2, figsize=figsize)
+        self._plot_image(
+            orig,
+            axes=axes[0],
+            figsize=figsize,
+            render_factor=render_factor,
+            display_render_factor=False,
+        )
+        self._plot_image(
+            result,
+            axes=axes[1],
+            figsize=figsize,
+            render_factor=render_factor,
+            display_render_factor=display_render_factor,
+        )
+
+    def _plot_solo(
+        self,
+        figsize: Tuple[int, int],
+        render_factor: int,
+        display_render_factor: bool,
+        result: Image,
+    ):
+        fig, axes = plt.subplots(1, 1, figsize=figsize)
+        self._plot_image(
+            result,
+            axes=axes,
+            figsize=figsize,
+            render_factor=render_factor,
+            display_render_factor=display_render_factor,
+        )
+
+    def _save_result_image(self, source_path: Path, image: Image, results_dir = None) -> Path:
+        if results_dir is None:
+            results_dir = Path(self.results_dir)
+        result_path = results_dir / source_path.name
+        image.save(result_path)
+        return result_path
+
+    def get_transformed_image(
+        self, path: Path, render_factor: int = None, post_process: bool = True,
+        watermarked: bool = True,
+    ) -> Image:
+        self._clean_mem()
+        orig_image = self._open_pil_image(path)
+        filtered_image = self.filter.filter(
+            orig_image, orig_image, render_factor=render_factor,post_process=post_process
+        )
+
+        if watermarked:
+            return get_watermarked(filtered_image)
+
+        return filtered_image
+
+    def _plot_image(
+        self,
+        image: Image,
+        render_factor: int,
+        axes: Axes = None,
+        figsize=(20, 20),
+        display_render_factor = False,
+    ):
+        if axes is None:
+            _, axes = plt.subplots(figsize=figsize)
+        axes.imshow(np.asarray(image) / 255)
+        axes.axis('off')
+        if render_factor is not None and display_render_factor:
+            plt.text(
+                10,
+                10,
+                'render_factor: ' + str(render_factor),
+                color='white',
+                backgroundcolor='black',
+            )
+
+    def _get_num_rows_columns(self, num_images: int, max_columns: int) -> Tuple[int, int]:
+        columns = min(num_images, max_columns)
+        rows = num_images // columns
+        rows = rows if rows * columns == num_images else rows + 1
+        return rows, columns
+
+
+class VideoColorizer:
+    def __init__(self, vis: ModelImageVisualizer):
+        self.vis = vis
+        workfolder = Path('./video')
+        self.source_folder = workfolder / "source"
+        self.bwframes_root = workfolder / "bwframes"
+        self.audio_root = workfolder / "audio"
+        self.colorframes_root = workfolder / "colorframes"
+        self.result_folder = workfolder / "result"
+
+    def _purge_images(self, dir):
+        for f in os.listdir(dir):
+            if re.search('.*?\.jpg', f):
+                os.remove(os.path.join(dir, f))
+
+    def _get_ffmpeg_probe(self, path:Path):
+        try:
+            probe = ffmpeg.probe(str(path))
+            return probe
+        except ffmpeg.Error as e:
+            logging.error("ffmpeg error: {0}".format(e), exc_info=True)
+            logging.error('stdout:' + e.stdout.decode('UTF-8'))
+            logging.error('stderr:' + e.stderr.decode('UTF-8'))
+            raise e
+        except Exception as e:
+            logging.error('Failed to instantiate ffmpeg.probe.  Details: {0}'.format(e), exc_info=True)   
+            raise e
+
+    def _get_fps(self, source_path: Path) -> str:
+        probe = self._get_ffmpeg_probe(source_path)
+        stream_data = next(
+            (stream for stream in probe['streams'] if stream['codec_type'] == 'video'),
+            None,
+        )
+        return stream_data['avg_frame_rate']
+
+    def _download_video_from_url(self, source_url, source_path: Path):
+        if source_path.exists():
+            source_path.unlink()
+
+        ydl_opts = {
+            'format': 'bestvideo[ext=mp4]+bestaudio[ext=m4a]/mp4',
+            'outtmpl': str(source_path),
+            'retries': 30,
+            'fragment-retries': 30
+        }
+        with youtube_dl.YoutubeDL(ydl_opts) as ydl:
+            ydl.download([source_url])
+
+    def _extract_raw_frames(self, source_path: Path):
+        bwframes_folder = self.bwframes_root / (source_path.stem)
+        bwframe_path_template = str(bwframes_folder / '%5d.jpg')
+        bwframes_folder.mkdir(parents=True, exist_ok=True)
+        self._purge_images(bwframes_folder)
+
+        process = (
+            ffmpeg
+                .input(str(source_path))
+                .output(str(bwframe_path_template), format='image2', vcodec='mjpeg', **{'q:v':'0'})
+                .global_args('-hide_banner')
+                .global_args('-nostats')
+                .global_args('-loglevel', 'error')
+        )
+
+        try:
+            process.run()
+        except ffmpeg.Error as e:
+            logging.error("ffmpeg error: {0}".format(e), exc_info=True)
+            logging.error('stdout:' + e.stdout.decode('UTF-8'))
+            logging.error('stderr:' + e.stderr.decode('UTF-8'))
+            raise e
+        except Exception as e:
+            logging.error('Errror while extracting raw frames from source video.  Details: {0}'.format(e), exc_info=True)   
+            raise e
+
+    def _colorize_raw_frames(
+        self, source_path: Path, render_factor: int = None, post_process: bool = True,
+        watermarked: bool = True,
+    ):
+        colorframes_folder = self.colorframes_root / (source_path.stem)
+        colorframes_folder.mkdir(parents=True, exist_ok=True)
+        self._purge_images(colorframes_folder)
+        bwframes_folder = self.bwframes_root / (source_path.stem)
+
+        for img in progress_bar(os.listdir(str(bwframes_folder))):
+            img_path = bwframes_folder / img
+
+            if os.path.isfile(str(img_path)):
+                color_image = self.vis.get_transformed_image(
+                    str(img_path), render_factor=render_factor, post_process=post_process,watermarked=watermarked
+                )
+                color_image.save(str(colorframes_folder / img))
+
+    def _build_video(self, source_path: Path) -> Path:
+        colorized_path = self.result_folder / (
+            source_path.name.replace('.mp4', '_no_audio.mp4')
+        )
+        colorframes_folder = self.colorframes_root / (source_path.stem)
+        colorframes_path_template = str(colorframes_folder / '%5d.jpg')
+        colorized_path.parent.mkdir(parents=True, exist_ok=True)
+        if colorized_path.exists():
+            colorized_path.unlink()
+        fps = self._get_fps(source_path)
+
+        process = (
+            ffmpeg 
+                .input(str(colorframes_path_template), format='image2', vcodec='mjpeg', framerate=fps) 
+                .output(str(colorized_path), crf=17, vcodec='libx264')
+                .global_args('-hide_banner')
+                .global_args('-nostats')
+                .global_args('-loglevel', 'error')
+        )
+
+        try:
+            process.run()
+        except ffmpeg.Error as e:
+            logging.error("ffmpeg error: {0}".format(e), exc_info=True)
+            logging.error('stdout:' + e.stdout.decode('UTF-8'))
+            logging.error('stderr:' + e.stderr.decode('UTF-8'))
+            raise e
+        except Exception as e:
+            logging.error('Errror while building output video.  Details: {0}'.format(e), exc_info=True)   
+            raise e
+
+        result_path = self.result_folder / source_path.name
+        if result_path.exists():
+            result_path.unlink()
+        # making copy of non-audio version in case adding back audio doesn't apply or fails.
+        shutil.copyfile(str(colorized_path), str(result_path))
+
+        # adding back sound here
+        audio_file = Path(str(source_path).replace('.mp4', '.aac'))
+        if audio_file.exists():
+            audio_file.unlink()
+
+        os.system(
+            'ffmpeg -y -i "'
+            + str(source_path)
+            + '" -vn -acodec copy "'
+            + str(audio_file)
+            + '"'
+            + ' -hide_banner'
+            + ' -nostats'
+            + ' -loglevel error'
+        )
+
+        if audio_file.exists():
+            os.system(
+                'ffmpeg -y -i "'
+                + str(colorized_path)
+                + '" -i "'
+                + str(audio_file)
+                + '" -shortest -c:v copy -c:a aac -b:a 256k "'
+                + str(result_path)
+                + '"'
+                + ' -hide_banner'
+                + ' -nostats'
+                + ' -loglevel error'
+            )
+        logging.info('Video created here: ' + str(result_path))
+        return result_path
+
+    def colorize_from_url(
+        self,
+        source_url,
+        file_name: str,
+        render_factor: int = None,
+        post_process: bool = True,
+        watermarked: bool = True,
+
+    ) -> Path:
+        source_path = self.source_folder / file_name
+        self._download_video_from_url(source_url, source_path)
+        return self._colorize_from_path(
+            source_path, render_factor=render_factor, post_process=post_process,watermarked=watermarked
+        )
+
+    def colorize_from_file_name(
+        self, file_name: str, render_factor: int = None,  watermarked: bool = True, post_process: bool = True,
+    ) -> Path:
+        source_path = self.source_folder / file_name
+        return self._colorize_from_path(
+            source_path, render_factor=render_factor,  post_process=post_process,watermarked=watermarked
+        )
+
+    def _colorize_from_path(
+        self, source_path: Path, render_factor: int = None,  watermarked: bool = True, post_process: bool = True
+    ) -> Path:
+        if not source_path.exists():
+            raise Exception(
+                'Video at path specfied, ' + str(source_path) + ' could not be found.'
+            )
+        self._extract_raw_frames(source_path)
+        self._colorize_raw_frames(
+            source_path, render_factor=render_factor,post_process=post_process,watermarked=watermarked
+        )
+        return self._build_video(source_path)
+
+
+def get_video_colorizer(render_factor: int = 21) -> VideoColorizer:
+    return get_stable_video_colorizer(render_factor=render_factor)
+
+
+def get_artistic_video_colorizer(
+    root_folder: Path = Path('./'),
+    weights_name: str = 'ColorizeArtistic_gen',
+    results_dir='result_images',
+    render_factor: int = 35
+) -> VideoColorizer:
+    learn = gen_inference_deep(root_folder=root_folder, weights_name=weights_name)
+    filtr = MasterFilter([ColorizerFilter(learn=learn)], render_factor=render_factor)
+    vis = ModelImageVisualizer(filtr, results_dir=results_dir)
+    return VideoColorizer(vis)
+
+
+def get_stable_video_colorizer(
+    root_folder: Path = Path('./'),
+    weights_name: str = 'ColorizeVideo_gen',
+    results_dir='result_images',
+    render_factor: int = 21
+) -> VideoColorizer:
+    learn = gen_inference_wide(root_folder=root_folder, weights_name=weights_name)
+    filtr = MasterFilter([ColorizerFilter(learn=learn)], render_factor=render_factor)
+    vis = ModelImageVisualizer(filtr, results_dir=results_dir)
+    return VideoColorizer(vis)
+
+
+def get_image_colorizer(
+    root_folder: Path = Path('./'), render_factor: int = 35, artistic: bool = True
+) -> ModelImageVisualizer:
+    if artistic:
+        return get_artistic_image_colorizer(root_folder=root_folder, render_factor=render_factor)
+    else:
+        return get_stable_image_colorizer(root_folder=root_folder, render_factor=render_factor)
+
+
+def get_stable_image_colorizer(
+    root_folder: Path = Path('./'),
+    weights_name: str = 'ColorizeStable_gen',
+    results_dir='result_images',
+    render_factor: int = 35
+) -> ModelImageVisualizer:
+    learn = gen_inference_wide(root_folder=root_folder, weights_name=weights_name)
+    filtr = MasterFilter([ColorizerFilter(learn=learn)], render_factor=render_factor)
+    vis = ModelImageVisualizer(filtr, results_dir=results_dir)
+    return vis
+
+
+def get_artistic_image_colorizer(
+    root_folder: Path = Path('./'),
+    weights_name: str = 'ColorizeArtistic_gen',
+    results_dir='result_images',
+    render_factor: int = 35
+) -> ModelImageVisualizer:
+    learn = gen_inference_deep(root_folder=root_folder, weights_name=weights_name)
+    filtr = MasterFilter([ColorizerFilter(learn=learn)], render_factor=render_factor)
+    vis = ModelImageVisualizer(filtr, results_dir=results_dir)
+    return vis
+
+
+def show_image_in_notebook(image_path: Path):
+    ipythondisplay.display(ipythonimage(str(image_path)))
+
+
+def show_video_in_notebook(video_path: Path):
+    video = io.open(video_path, 'r+b').read()
+    encoded = base64.b64encode(video)
+    ipythondisplay.display(
+        HTML(
+            data='''<video alt="test" autoplay 
+                loop controls style="height: 400px;">
+                <source src="data:video/mp4;base64,{0}" type="video/mp4" />
+             </video>'''.format(
+                encoded.decode('ascii')
+            )
+        )
+    )

DeOldify/environment.yml

@@ -0,0 +1,23 @@
+name: deoldify
+channels:
+- fastai
+- conda-forge
+- defaults
+dependencies:
+- pip
+- fastai=1.0.60
+- mkl=2024.0
+- python=3.10
+- pytorch::pytorch=1.11.0
+- pytorch::torchvision
+- pytorch::torchaudio
+- tensorboardX
+- jupyterlab
+- pillow>=9.0.0
+- ipywidgets
+- ffmpeg
+- pip:
+  - ffmpeg-python
+  - opencv-python>=4.2.0.32
+  - wandb
+  - yt-dlp

DeOldify/fastai/LICENSE

@@ -0,0 +1,60 @@
+Apache License, Version 2.0 Apache License Version 2.0, January 2004 http://www.apache.org/licenses/
+
+TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+1. Definitions.
+
+"License" shall mean the terms and conditions for use, reproduction, and distribution as defined by Sections 1 through 9 of this document.
+
+"Licensor" shall mean the copyright owner or entity authorized by the copyright owner that is granting the License.
+
+"Legal Entity" shall mean the union of the acting entity and all other entities that control, are controlled by, or are under common control with that entity. For the purposes of this definition, "control" means (i) the power, direct or indirect, to cause the direction or management of such entity, whether by contract or otherwise, or (ii) ownership of fifty percent (50%) or more of the outstanding shares, or (iii) beneficial ownership of such entity.
+
+"You" (or "Your") shall mean an individual or Legal Entity exercising permissions granted by this License.
+
+"Source" form shall mean the preferred form for making modifications, including but not limited to software source code, documentation source, and configuration files.
+
+"Object" form shall mean any form resulting from mechanical transformation or translation of a Source form, including but not limited to compiled object code, generated documentation, and conversions to other media types.
+
+"Work" shall mean the work of authorship, whether in Source or Object form, made available under the License, as indicated by a copyright notice that is included in or attached to the work (an example is provided in the Appendix below).
+
+"Derivative Works" shall mean any work, whether in Source or Object form, that is based on (or derived from) the Work and for which the editorial revisions, annotations, elaborations, or other modifications represent, as a whole, an original work of authorship. For the purposes of this License, Derivative Works shall not include works that remain separable from, or merely link (or bind by name) to the interfaces of, the Work and Derivative Works thereof.
+
+"Contribution" shall mean any work of authorship, including the original version of the Work and any modifications or additions to that Work or Derivative Works thereof, that is intentionally submitted to Licensor for inclusion in the Work by the copyright owner or by an individual or Legal Entity authorized to submit on behalf of the copyright owner. For the purposes of this definition, "submitted" means any form of electronic, verbal, or written communication sent to the Licensor or its representatives, including but not limited to communication on electronic mailing lists, source code control systems, and issue tracking systems that are managed by, or on behalf of, the Licensor for the purpose of discussing and improving the Work, but excluding communication that is conspicuously marked or otherwise designated in writing by the copyright owner as "Not a Contribution."
+
+"Contributor" shall mean Licensor and any individual or Legal Entity on behalf of whom a Contribution has been received by Licensor and subsequently incorporated within the Work.
+
+2. Grant of Copyright License.
+
+Subject to the terms and conditions of this License, each Contributor hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable copyright license to reproduce, prepare Derivative Works of, publicly display, publicly perform, sublicense, and distribute the Work and such Derivative Works in Source or Object form.
+
+3. Grant of Patent License.
+
+Subject to the terms and conditions of this License, each Contributor hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable (except as stated in this section) patent license to make, have made, use, offer to sell, sell, import, and otherwise transfer the Work, where such license applies only to those patent claims licensable by such Contributor that are necessarily infringed by their Contribution(s) alone or by combination of their Contribution(s) with the Work to which such Contribution(s) was submitted. If You institute patent litigation against any entity (including a cross-claim or counterclaim in a lawsuit) alleging that the Work or a Contribution incorporated within the Work constitutes direct or contributory patent infringement, then any patent licenses granted to You under this License for that Work shall terminate as of the date such litigation is filed.
+
+4. Redistribution.
+
+You may reproduce and distribute copies of the Work or Derivative Works thereof in any medium, with or without modifications, and in Source or Object form, provided that You meet the following conditions:
+
+You must give any other recipients of the Work or Derivative Works a copy of this License; and You must cause any modified files to carry prominent notices stating that You changed the files; and You must retain, in the Source form of any Derivative Works that You distribute, all copyright, patent, trademark, and attribution notices from the Source form of the Work, excluding those notices that do not pertain to any part of the Derivative Works; and If the Work includes a "NOTICE" text file as part of its distribution, then any Derivative Works that You distribute must include a readable copy of the attribution notices contained within such NOTICE file, excluding those notices that do not pertain to any part of the Derivative Works, in at least one of the following places: within a NOTICE text file distributed as part of the Derivative Works; within the Source form or documentation, if provided along with the Derivative Works; or, within a display generated by the Derivative Works, if and wherever such third-party notices normally appear. The contents of the NOTICE file are for informational purposes only and do not modify the License. You may add Your own attribution notices within Derivative Works that You distribute, alongside or as an addendum to the NOTICE text from the Work, provided that such additional attribution notices cannot be construed as modifying the License. You may add Your own copyright statement to Your modifications and may provide additional or different license terms and conditions for use, reproduction, or distribution of Your modifications, or for any such Derivative Works as a whole, provided Your use, reproduction, and distribution of the Work otherwise complies with the conditions stated in this License.
+
+5. Submission of Contributions.
+
+Unless You explicitly state otherwise, any Contribution intentionally submitted for inclusion in the Work by You to the Licensor shall be under the terms and conditions of this License, without any additional terms or conditions. Notwithstanding the above, nothing herein shall supersede or modify the terms of any separate license agreement you may have executed with Licensor regarding such Contributions.
+
+6. Trademarks.
+
+This License does not grant permission to use the trade names, trademarks, service marks, or product names of the Licensor, except as required for reasonable and customary use in describing the origin of the Work and reproducing the content of the NOTICE file.
+
+7. Disclaimer of Warranty.
+
+Unless required by applicable law or agreed to in writing, Licensor provides the Work (and each Contributor provides its Contributions) on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied, including, without limitation, any warranties or conditions of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A PARTICULAR PURPOSE. You are solely responsible for determining the appropriateness of using or redistributing the Work and assume any risks associated with Your exercise of permissions under this License.
+
+8. Limitation of Liability.
+
+In no event and under no legal theory, whether in tort (including negligence), contract, or otherwise, unless required by applicable law (such as deliberate and grossly negligent acts) or agreed to in writing, shall any Contributor be liable to You for damages, including any direct, indirect, special, incidental, or consequential damages of any character arising as a result of this License or out of the use or inability to use the Work (including but not limited to damages for loss of goodwill, work stoppage, computer failure or malfunction, or any and all other commercial damages or losses), even if such Contributor has been advised of the possibility of such damages.
+
+9. Accepting Warranty or Additional Liability.
+
+While redistributing the Work or Derivative Works thereof, You may choose to offer, and charge a fee for, acceptance of support, warranty, indemnity, or other liability obligations and/or rights consistent with this License. However, in accepting such obligations, You may act only on Your own behalf and on Your sole responsibility, not on behalf of any other Contributor, and only if You agree to indemnify, defend, and hold each Contributor harmless for any liability incurred by, or claims asserted against, such Contributor by reason of your accepting any such warranty or additional liability.
+

DeOldify/fastai/__init__.py

@@ -0,0 +1,2 @@
+from .version import __version__
+

DeOldify/fastai/basic_data.py

@@ -0,0 +1,279 @@
+"`fastai.data` loads and manages datasets with `DataBunch`"
+from .torch_core import *
+from torch.utils.data.dataloader import default_collate
+
+DatasetType = Enum('DatasetType', 'Train Valid Test Single Fix')
+__all__ = ['DataBunch', 'DeviceDataLoader', 'DatasetType', 'load_data']
+
+old_dl_init = torch.utils.data.DataLoader.__init__
+
+def intercept_args(self, dataset, batch_size=1, shuffle=False, sampler=None, batch_sampler=None,
+                 num_workers=0, collate_fn=default_collate, pin_memory=True, drop_last=False,
+                 timeout=0, worker_init_fn=None):
+    self.init_kwargs = {'batch_size':batch_size, 'shuffle':shuffle, 'sampler':sampler, 'batch_sampler':batch_sampler,
+                        'num_workers':num_workers, 'collate_fn':collate_fn, 'pin_memory':pin_memory,
+                        'drop_last': drop_last, 'timeout':timeout, 'worker_init_fn':worker_init_fn}
+    old_dl_init(self, dataset, **self.init_kwargs)
+
+torch.utils.data.DataLoader.__init__ = intercept_args
+
+def DataLoader___getattr__(dl, k:str)->Any: return getattr(dl.dataset, k)
+DataLoader.__getattr__ = DataLoader___getattr__
+
+def DataLoader___setstate__(dl, data:Any): dl.__dict__.update(data)
+DataLoader.__setstate__ = DataLoader___setstate__
+
+@dataclass
+class DeviceDataLoader():
+    "Bind a `DataLoader` to a `torch.device`."
+    dl: DataLoader
+    device: torch.device
+    tfms: List[Callable]=None
+    collate_fn: Callable=data_collate
+    def __post_init__(self):
+        self.dl.collate_fn=self.collate_fn
+        self.tfms = listify(self.tfms)
+
+    def __len__(self)->int: return len(self.dl)
+    def __getattr__(self,k:str)->Any: return getattr(self.dl, k)
+    def __setstate__(self,data:Any): self.__dict__.update(data)
+
+    @property
+    def batch_size(self):   return self.dl.batch_size
+    @batch_size.setter
+    def batch_size(self,v):
+        new_kwargs = {**self.dl.init_kwargs, 'batch_size':v, 'collate_fn':self.collate_fn}
+        self.dl = self.dl.__class__(self.dl.dataset, **new_kwargs)
+        if hasattr(self.dl.dataset, 'bs'): self.dl.dataset.bs = v
+
+    @property
+    def num_workers(self):   return self.dl.num_workers
+    @num_workers.setter
+    def num_workers(self,v): self.dl.num_workers = v
+
+    def add_tfm(self,tfm:Callable)->None:
+        "Add `tfm` to `self.tfms`."
+        self.tfms.append(tfm)
+    def remove_tfm(self,tfm:Callable)->None:
+        "Remove `tfm` from `self.tfms`."
+        if tfm in self.tfms: self.tfms.remove(tfm)
+
+    def new(self, **kwargs):
+        "Create a new copy of `self` with `kwargs` replacing current values."
+        new_kwargs = {**self.dl.init_kwargs, **kwargs}
+        return DeviceDataLoader(self.dl.__class__(self.dl.dataset, **new_kwargs), self.device, self.tfms,
+                                self.collate_fn)
+
+    def proc_batch(self,b:Tensor)->Tensor:
+        "Process batch `b` of `TensorImage`."
+        b = to_device(b, self.device)
+        for f in listify(self.tfms): b = f(b)
+        return b
+
+    def __iter__(self):
+        "Process and returns items from `DataLoader`."
+        for b in self.dl: yield self.proc_batch(b)
+
+    @classmethod
+    def create(cls, dataset:Dataset, bs:int=64, shuffle:bool=False, device:torch.device=defaults.device,
+               tfms:Collection[Callable]=tfms, num_workers:int=defaults.cpus, collate_fn:Callable=data_collate, **kwargs:Any):
+        "Create DeviceDataLoader from `dataset` with `bs` and `shuffle`: process using `num_workers`."
+        return cls(DataLoader(dataset, batch_size=bs, shuffle=shuffle, num_workers=num_workers, **kwargs),
+                   device=device, tfms=tfms, collate_fn=collate_fn)
+
+class DataBunch():
+    "Bind `train_dl`,`valid_dl` and `test_dl` in a data object."
+
+    def __init__(self, train_dl:DataLoader, valid_dl:DataLoader, fix_dl:DataLoader=None, test_dl:Optional[DataLoader]=None,
+                 device:torch.device=None, dl_tfms:Optional[Collection[Callable]]=None, path:PathOrStr='.',
+                 collate_fn:Callable=data_collate, no_check:bool=False):
+        self.dl_tfms = listify(dl_tfms)
+        self.device = defaults.device if device is None else device
+        assert not isinstance(train_dl,DeviceDataLoader)
+        def _create_dl(dl, **kwargs):
+            if dl is None: return None
+            return DeviceDataLoader(dl, self.device, self.dl_tfms, collate_fn, **kwargs)
+        self.train_dl,self.valid_dl,self.fix_dl,self.test_dl = map(_create_dl, [train_dl,valid_dl,fix_dl,test_dl])
+        if fix_dl is None: self.fix_dl = self.train_dl.new(shuffle=False, drop_last=False)
+        self.single_dl = _create_dl(DataLoader(valid_dl.dataset, batch_size=1, num_workers=0))
+        self.path = Path(path)
+        if not no_check: self.sanity_check()
+
+    def __repr__(self)->str:
+        return f'{self.__class__.__name__};\n\nTrain: {self.train_ds};\n\nValid: {self.valid_ds};\n\nTest: {self.test_ds}'
+
+    @staticmethod
+    def _init_ds(train_ds:Dataset, valid_ds:Dataset, test_ds:Optional[Dataset]=None):
+        # train_ds, but without training tfms
+        fix_ds = valid_ds.new(train_ds.x, train_ds.y) if hasattr(valid_ds,'new') else train_ds
+        return [o for o in (train_ds,valid_ds,fix_ds,test_ds) if o is not None]
+
+    @classmethod
+    def create(cls, train_ds:Dataset, valid_ds:Dataset, test_ds:Optional[Dataset]=None, path:PathOrStr='.', bs:int=64,
+               val_bs:int=None, num_workers:int=defaults.cpus, dl_tfms:Optional[Collection[Callable]]=None,
+               device:torch.device=None, collate_fn:Callable=data_collate, no_check:bool=False, **dl_kwargs)->'DataBunch':
+        "Create a `DataBunch` from `train_ds`, `valid_ds` and maybe `test_ds` with a batch size of `bs`. Passes `**dl_kwargs` to `DataLoader()`"
+        datasets = cls._init_ds(train_ds, valid_ds, test_ds)
+        val_bs = ifnone(val_bs, bs)
+        dls = [DataLoader(d, b, shuffle=s, drop_last=s, num_workers=num_workers, **dl_kwargs) for d,b,s in
+               zip(datasets, (bs,val_bs,val_bs,val_bs), (True,False,False,False)) if d is not None]
+        return cls(*dls, path=path, device=device, dl_tfms=dl_tfms, collate_fn=collate_fn, no_check=no_check)
+
+    def __getattr__(self,k:int)->Any: return getattr(self.train_dl, k)
+    def __setstate__(self,data:Any): self.__dict__.update(data)
+
+    def dl(self, ds_type:DatasetType=DatasetType.Valid)->DeviceDataLoader:
+        "Returns appropriate `Dataset` for validation, training, or test (`ds_type`)."
+        #TODO: refactor
+        return (self.train_dl if ds_type == DatasetType.Train else
+                self.test_dl if ds_type == DatasetType.Test else
+                self.valid_dl if ds_type == DatasetType.Valid else
+                self.single_dl if ds_type == DatasetType.Single else
+                self.fix_dl)
+
+    @property
+    def dls(self)->List[DeviceDataLoader]:
+        "Returns a list of all DeviceDataLoaders. If you need a specific DeviceDataLoader, access via the relevant property (`train_dl`, `valid_dl`, etc) as the index of DLs in this list is not guaranteed to remain constant."
+        res = [self.train_dl, self.fix_dl, self.single_dl]
+        # Preserve the original ordering of Train, Valid, Fix, Single, Test Data Loaders
+        # (Unknown/not verified as of 1.0.47 whether there are other methods explicitly using DLs their list index)
+        if self.valid_dl: res.insert(1, self.valid_dl)
+        return res if not self.test_dl else res + [self.test_dl]
+
+    def add_tfm(self,tfm:Callable)->None:
+        for dl in self.dls: dl.add_tfm(tfm)
+
+    def remove_tfm(self,tfm:Callable)->None:
+        for dl in self.dls: dl.remove_tfm(tfm)
+
+    def save(self, file:PathLikeOrBinaryStream= 'data_save.pkl')->None:
+        "Save the `DataBunch` in `self.path/file`. `file` can be file-like (file or buffer)"
+        if not getattr(self, 'label_list', False):
+            warn("Serializing the `DataBunch` only works when you created it using the data block API.")
+            return
+        try_save(self.label_list, self.path, file)
+
+    def add_test(self, items:Iterator, label:Any=None, tfms=None, tfm_y=None)->None:
+        "Add the `items` as a test set. Pass along `label` otherwise label them with `EmptyLabel`."
+        self.label_list.add_test(items, label=label, tfms=tfms, tfm_y=tfm_y)
+        vdl = self.valid_dl
+        dl = DataLoader(self.label_list.test, vdl.batch_size, shuffle=False, drop_last=False, num_workers=vdl.num_workers)
+        self.test_dl = DeviceDataLoader(dl, vdl.device, vdl.tfms, vdl.collate_fn)
+
+    def one_batch(self, ds_type:DatasetType=DatasetType.Train, detach:bool=True, denorm:bool=True, cpu:bool=True)->Collection[Tensor]:
+        "Get one batch from the data loader of `ds_type`. Optionally `detach` and `denorm`."
+        dl = self.dl(ds_type)
+        w = self.num_workers
+        self.num_workers = 0
+        try:     x,y = next(iter(dl))
+        finally: self.num_workers = w
+        if detach: x,y = to_detach(x,cpu=cpu),to_detach(y,cpu=cpu)
+        norm = getattr(self,'norm',False)
+        if denorm and norm:
+            x = self.denorm(x)
+            if norm.keywords.get('do_y',False): y = self.denorm(y, do_x=True)
+        return x,y
+
+    def one_item(self, item, detach:bool=False, denorm:bool=False, cpu:bool=False):
+        "Get `item` into a batch. Optionally `detach` and `denorm`."
+        ds = self.single_ds
+        with ds.set_item(item):
+            return self.one_batch(ds_type=DatasetType.Single, detach=detach, denorm=denorm, cpu=cpu)
+
+    def show_batch(self, rows:int=5, ds_type:DatasetType=DatasetType.Train, reverse:bool=False, **kwargs)->None:
+        "Show a batch of data in `ds_type` on a few `rows`."
+        x,y = self.one_batch(ds_type, True, True)
+        if reverse: x,y = x.flip(0),y.flip(0)
+        n_items = rows **2 if self.train_ds.x._square_show else rows
+        if self.dl(ds_type).batch_size < n_items: n_items = self.dl(ds_type).batch_size
+        xs = [self.train_ds.x.reconstruct(grab_idx(x, i)) for i in range(n_items)]
+        #TODO: get rid of has_arg if possible
+        if has_arg(self.train_ds.y.reconstruct, 'x'):
+            ys = [self.train_ds.y.reconstruct(grab_idx(y, i), x=x) for i,x in enumerate(xs)]
+        else : ys = [self.train_ds.y.reconstruct(grab_idx(y, i)) for i in range(n_items)]
+        self.train_ds.x.show_xys(xs, ys, **kwargs)
+ 
+    def export(self, file:PathLikeOrBinaryStream='export.pkl'):
+        "Export the minimal state of `self` for inference in `self.path/file`. `file` can be file-like (file or buffer)"
+        xtra = dict(normalize=self.norm.keywords) if getattr(self, 'norm', False) else {}
+        try_save(self.valid_ds.get_state(**xtra), self.path, file)
+
+    def _grab_dataset(self, dl:DataLoader):
+        ds = dl.dl.dataset
+        while hasattr(ds, 'dataset'): ds = ds.dataset
+        return ds
+
+    @property
+    def train_ds(self)->Dataset: return self._grab_dataset(self.train_dl)
+    @property
+    def valid_ds(self)->Dataset: return self._grab_dataset(self.valid_dl)
+    @property
+    def single_ds(self)->Dataset: return self._grab_dataset(self.single_dl)
+    @property
+    def loss_func(self)->OptLossFunc:
+        return getattr(self.train_ds.y, 'loss_func', F.nll_loss) if hasattr(self.train_ds, 'y') else F.nll_loss
+
+    @property
+    def test_ds(self)->Dataset:
+        return self._grab_dataset(self.test_dl) if self.test_dl is not None else None
+
+    @property
+    def empty_val(self)->bool:
+        if not hasattr(self, 'valid_dl') or self.valid_dl is None:            return True
+        if hasattr(self.valid_ds, 'items') and len(self.valid_ds.items) == 0: return True
+        return (len(self.valid_ds) == 0)
+
+    @property
+    def is_empty(self)->bool:
+        return not ((self.train_dl and len(self.train_ds.items) != 0) or 
+                    (self.valid_dl and len(self.valid_ds.items) != 0) or 
+                    (self.test_dl  and len(self.test_ds.items)  != 0))
+    
+    @property
+    def batch_size(self):   return self.train_dl.batch_size
+    @batch_size.setter
+    def batch_size(self,v):
+        self.train_dl.batch_size,self.valid_dl.batch_size = v,v
+        if self.test_dl is not None: self.test_dl.batch_size = v
+
+    def sanity_check(self):
+        "Check the underlying data in the training set can be properly loaded."
+        final_message = "You can deactivate this warning by passing `no_check=True`."
+        if not hasattr(self.train_ds, 'items') or len(self.train_ds.items) == 0 or not hasattr(self.train_dl, 'batch_sampler'): return
+        if len(self.train_dl) == 0:
+            warn(f"""Your training dataloader is empty, you have only {len(self.train_dl.dataset)} items in your training set.
+                 Your batch size is {self.train_dl.batch_size}, you should lower it.""")
+            print(final_message)
+            return
+        idx = next(iter(self.train_dl.batch_sampler))
+        samples,fails = [],[]
+        for i in idx:
+            try:    samples.append(self.train_dl.dataset[i])
+            except: fails.append(i)
+        if len(fails) > 0:
+            warn_msg = "There seems to be something wrong with your dataset, for example, in the first batch can't access"
+            if len(fails) == len(idx):
+                warn_msg += f" any element of self.train_ds.\nTried: {show_some(idx)}"
+            else:
+                warn_msg += f" these elements in self.train_ds: {show_some(fails)}"
+            warn(warn_msg)
+            print(final_message)
+            return
+        try: batch = self.collate_fn(samples)
+        except:
+            message = "It's not possible to collate samples of your dataset together in a batch."
+            try:
+                shapes = [[o[i].data.shape for o in samples] for i in range(2)]
+                message += f'\nShapes of the inputs/targets:\n{shapes}'
+            except: pass
+            warn(message)
+            print(final_message)
+
+def load_data(path:PathOrStr, file:PathLikeOrBinaryStream='data_save.pkl', bs:int=64, val_bs:int=None, num_workers:int=defaults.cpus,
+              dl_tfms:Optional[Collection[Callable]]=None, device:torch.device=None, collate_fn:Callable=data_collate,
+              no_check:bool=False, **kwargs)->DataBunch:
+    "Load a saved `DataBunch` from `path/file`. `file` can be file-like (file or buffer)"
+    source = Path(path)/file if is_pathlike(file) else file
+    ll = torch.load(source, map_location='cpu') if defaults.device == torch.device('cpu') else torch.load(source)
+    return ll.databunch(path=path, bs=bs, val_bs=val_bs, num_workers=num_workers, dl_tfms=dl_tfms, device=device,
+                        collate_fn=collate_fn, no_check=no_check, **kwargs)

DeOldify/fastai/basic_train.py

@@ -0,0 +1,629 @@
+"Provides basic training and validation with `Learner`"
+from .torch_core import *
+from .basic_data import *
+from .callback import *
+from .data_block import *
+from .utils.ipython import gpu_mem_restore
+import inspect
+from fastprogress.fastprogress import format_time, IN_NOTEBOOK
+from time import time
+from fastai.sixel import plot_sixel
+
+__all__ = ['Learner', 'LearnerCallback', 'Recorder', 'RecordOnCPU', 'fit', 'loss_batch', 'train_epoch', 'validate',
+           'get_preds', 'load_learner']
+
+defaults.lr = slice(3e-3)
+defaults.wd = 1e-2
+defaults.extra_callbacks    = None
+defaults.extra_callback_fns = None
+
+def loss_batch(model:nn.Module, xb:Tensor, yb:Tensor, loss_func:OptLossFunc=None, opt:OptOptimizer=None,
+               cb_handler:Optional[CallbackHandler]=None, count:[int]=[1], batch_multiplier:int=1)->Tuple[Union[Tensor,int,float,str]]:
+    "Calculate loss and metrics for a batch, call out to callbacks as necessary."
+    cb_handler = ifnone(cb_handler, CallbackHandler())
+    if not is_listy(xb): xb = [xb]
+    if not is_listy(yb): yb = [yb]
+    out = model(*xb)
+
+    if not loss_func: return to_detach(out), yb[0].detach()
+    out = cb_handler.on_loss_begin(out)
+    loss = loss_func(out, *yb)/batch_multiplier
+    count[0]-=1
+
+    if opt is not None:
+        loss,skip_bwd = cb_handler.on_backward_begin(loss)
+        if not skip_bwd:                     loss.backward()
+        if count[0] == 0:
+            if not cb_handler.on_backward_end(): opt.step()
+            if not cb_handler.on_step_end():     opt.zero_grad()
+            count[0] = batch_multiplier
+
+    return loss.detach().cpu()
+
+def get_preds(model:nn.Module, dl:DataLoader, pbar:Optional[PBar]=None, cb_handler:Optional[CallbackHandler]=None,
+              activ:nn.Module=None, loss_func:OptLossFunc=None, n_batch:Optional[int]=None) -> List[Tensor]:
+    "Tuple of predictions and targets, and optional losses (if `loss_func`) using `dl`, max batches `n_batch`."
+    res = [torch.cat(o).cpu() for o in
+           zip(*validate(model, dl, cb_handler=cb_handler, pbar=pbar, average=False, n_batch=n_batch))]
+    if loss_func is not None:
+        with NoneReduceOnCPU(loss_func) as lf: res.append(lf(res[0], res[1]))
+    if activ is not None: res[0] = activ(res[0])
+    return res
+
+def validate(model:nn.Module, dl:DataLoader, loss_func:OptLossFunc=None, cb_handler:Optional[CallbackHandler]=None,
+             pbar:Optional[PBar]=None, average=True, n_batch:Optional[int]=None)->Iterator[Tuple[Union[Tensor,int],...]]:
+    "Calculate `loss_func` of `model` on `dl` in evaluation mode."
+    model.eval()
+    with torch.no_grad():
+        val_losses,nums = [],[]
+        if cb_handler: cb_handler.set_dl(dl)
+        for xb,yb in progress_bar(dl, parent=pbar, leave=(pbar is not None)):
+            if cb_handler: xb, yb = cb_handler.on_batch_begin(xb, yb, train=False)
+            val_loss = loss_batch(model, xb, yb, loss_func, cb_handler=cb_handler)
+            val_losses.append(val_loss)
+            if not is_listy(yb): yb = [yb]
+            nums.append(first_el(yb).shape[0])
+            if cb_handler and cb_handler.on_batch_end(val_losses[-1]): break
+            if n_batch and (len(nums)>=n_batch): break
+        nums = np.array(nums, dtype=np.float32)
+        if average: return (to_np(torch.stack(val_losses)) * nums).sum() / nums.sum()
+        else:       return val_losses
+
+def train_epoch(model:nn.Module, dl:DataLoader, opt:optim.Optimizer, loss_func:LossFunction)->None:
+    "Simple training of `model` for 1 epoch of `dl` using optim `opt` and loss function `loss_func`."
+    model.train()
+    for xb,yb in dl:
+        loss = loss_func(model(xb), yb)
+        loss.backward()
+        opt.step()
+        opt.zero_grad()
+
+@dataclass
+class BasicLearner():
+    model:nn.Module
+    loss_func:LossFunction
+    opt:optim.Optimizer
+    data:DataBunch
+
+def fit(epochs:int, learn:BasicLearner, callbacks:Optional[CallbackList]=None, metrics:OptMetrics=None, batch_multiplier:int=1)->None:
+    "Fit the `model` on `data` and learn using `loss_func` and `opt`."
+    assert len(learn.data.train_dl) != 0, f"""Your training dataloader is empty, can't train a model.
+        Use a smaller batch size (batch size={learn.data.train_dl.batch_size} for {len(learn.data.train_dl.dataset)} elements)."""
+    cb_handler = CallbackHandler(callbacks, metrics)
+    pbar = master_bar(range(epochs))
+    cb_handler.on_train_begin(epochs, pbar=pbar, metrics=metrics)
+
+    exception=False
+    try:
+        for epoch in pbar:
+            learn.model.train()
+            cb_handler.set_dl(learn.data.train_dl)
+            cb_handler.on_epoch_begin()
+            count = [batch_multiplier]
+            for xb,yb in progress_bar(learn.data.train_dl, parent=pbar):
+                xb, yb = cb_handler.on_batch_begin(xb, yb)
+                loss = loss_batch(learn.model, xb, yb, learn.loss_func, learn.opt, cb_handler, count=count, batch_multiplier=batch_multiplier)
+                if cb_handler.on_batch_end(loss): break
+
+            if not cb_handler.skip_validate and not learn.data.empty_val:
+                val_loss = validate(learn.model, learn.data.valid_dl, loss_func=learn.loss_func,
+                                       cb_handler=cb_handler, pbar=pbar)
+            else: val_loss=None
+            if cb_handler.on_epoch_end(val_loss): break
+    except Exception as e:
+        exception = e
+        raise
+    finally: cb_handler.on_train_end(exception)
+
+loss_func_name2activ = {'cross_entropy_loss': F.softmax, 'nll_loss': torch.exp, 'poisson_nll_loss': torch.exp,
+    'kl_div_loss': torch.exp, 'bce_with_logits_loss': torch.sigmoid, 'cross_entropy': F.softmax,
+    'kl_div': torch.exp, 'binary_cross_entropy_with_logits': torch.sigmoid,
+}
+
+def _loss_func_name2activ(name:str, axis:int=-1):
+    res = loss_func_name2activ[name]
+    if res == F.softmax: res = partial(F.softmax, dim=axis)
+    return res
+
+def _loss_func2activ(loss_func):
+    if getattr(loss_func,'keywords',None):
+        if not loss_func.keywords.get('log_input', True): return
+    axis = getattr(loss_func, 'axis', -1)
+    # flattened loss
+    loss_func = getattr(loss_func, 'func', loss_func)
+    # could have a partial inside flattened loss! Duplicate on purpose.
+    loss_func = getattr(loss_func, 'func', loss_func)
+    cls_name = camel2snake(loss_func.__class__.__name__)
+    if cls_name == 'mix_up_loss':
+        loss_func = loss_func.crit
+        cls_name = camel2snake(loss_func.__class__.__name__)
+    if cls_name in loss_func_name2activ:
+        if cls_name == 'poisson_nll_loss' and (not getattr(loss_func, 'log_input', True)): return
+        return _loss_func_name2activ(cls_name, axis)
+    if getattr(loss_func,'__name__','') in loss_func_name2activ:
+        return _loss_func_name2activ(loss_func.__name__, axis)
+    return noop
+
+@dataclass
+class Learner():
+    "Trainer for `model` using `data` to minimize `loss_func` with optimizer `opt_func`."
+    data:DataBunch
+    model:nn.Module
+    opt_func:Callable=AdamW
+    loss_func:Callable=None
+    metrics:Collection[Callable]=None
+    true_wd:bool=True
+    bn_wd:bool=True
+    wd:Floats=defaults.wd
+    train_bn:bool=True
+    path:str = None
+    model_dir:PathOrStr = 'models'
+    callback_fns:Collection[Callable]=None
+    callbacks:Collection[Callback]=field(default_factory=list)
+    layer_groups:Collection[nn.Module]=None
+    add_time:bool=True
+    silent:bool=None
+    def __post_init__(self)->None:
+        "Setup path,metrics, callbacks and ensure model directory exists."
+        self.path = Path(ifnone(self.path, self.data.path))
+        self.model = self.model.to(self.data.device)
+        self.loss_func = self.loss_func or self.data.loss_func
+        self.metrics=listify(self.metrics)
+        if not self.layer_groups: self.layer_groups = [nn.Sequential(*flatten_model(self.model))]
+        self.callbacks = listify(self.callbacks)
+        if self.silent is None: self.silent = defaults.silent
+        self.callback_fns = [partial(Recorder, add_time=self.add_time, silent=self.silent)] + listify(self.callback_fns)
+
+    def init(self, init): apply_init(self.model, init)
+
+    def _test_writeable_path(self):
+        path = self.path/self.model_dir
+        try:
+            path.mkdir(parents=True, exist_ok=True)
+            tmp_file = get_tmp_file(path)
+        except OSError as e:
+            raise Exception(f"{e}\nCan't write to '{path}', set `learn.model_dir` attribute in Learner to a full libpath path that is writable") from None
+        os.remove(tmp_file)
+
+    def lr_range(self, lr:Union[float,slice])->np.ndarray:
+        "Build differential learning rates from `lr`."
+        if not isinstance(lr,slice): return lr
+        if lr.start: res = even_mults(lr.start, lr.stop, len(self.layer_groups))
+        else: res = [lr.stop/10]*(len(self.layer_groups)-1) + [lr.stop]
+        return np.array(res)
+
+    def fit(self, epochs:int, lr:Union[Floats,slice]=defaults.lr,
+            wd:Floats=None, callbacks:Collection[Callback]=None, batch_multiplier:int=1)->None:
+        "Fit the model on this learner with `lr` learning rate, `wd` weight decay for `epochs` with `callbacks`."
+        lr = self.lr_range(lr)
+        if wd is None: wd = self.wd
+        if not getattr(self, 'opt', False): self.create_opt(lr, wd)
+        else: self.opt.lr,self.opt.wd = lr,wd
+        callbacks = [cb(self) for cb in self.callback_fns + listify(defaults.extra_callback_fns)] + listify(callbacks)
+        if defaults.extra_callbacks is not None: callbacks += defaults.extra_callbacks
+        fit(epochs, self, metrics=self.metrics, callbacks=self.callbacks+callbacks, batch_multiplier=batch_multiplier)
+
+    def create_opt(self, lr:Floats, wd:Floats=0.)->None:
+        "Create optimizer with `lr` learning rate and `wd` weight decay."
+        self.opt = OptimWrapper.create(self.opt_func, lr, self.layer_groups, wd=wd, true_wd=self.true_wd, bn_wd=self.bn_wd)
+
+    def split(self, split_on:SplitFuncOrIdxList)->None:
+        "Split the model at `split_on`."
+        if isinstance(split_on,Callable): split_on = split_on(self.model)
+        self.layer_groups = split_model(self.model, split_on)
+        return self
+
+    def freeze_to(self, n:int)->None:
+        "Freeze layers up to layer group `n`."
+        for g in self.layer_groups[:n]:
+            for l in g:
+                if not self.train_bn or not isinstance(l, bn_types): requires_grad(l, False)
+        for g in self.layer_groups[n:]: requires_grad(g, True)
+        self.create_opt(defaults.lr)
+
+    def freeze(self)->None:
+        "Freeze up to last layer group."
+        assert(len(self.layer_groups)>1)
+        self.freeze_to(-1)
+
+    def unfreeze(self):
+        "Unfreeze entire model."
+        self.freeze_to(0)
+
+    def export(self, file:PathLikeOrBinaryStream='export.pkl', destroy=False):
+        "Export the state of the `Learner` in `self.path/file`. `file` can be file-like (file or buffer)"
+        if rank_distrib(): return # don't save if slave proc
+        args = ['opt_func', 'loss_func', 'metrics', 'true_wd', 'bn_wd', 'wd', 'train_bn', 'model_dir', 'callback_fns']
+        state = {a:getattr(self,a) for a in args}
+        state['cb_state'] = {cb.__class__:cb.get_state() for cb in self.callbacks}
+        #layer_groups -> need to find a way
+        #TO SEE: do we save model structure and weights separately?
+        with ModelOnCPU(self.model) as m:
+            state['model'] = m
+            xtra = dict(normalize=self.data.norm.keywords) if getattr(self.data, 'norm', False) else {}
+            state['data'] = self.data.valid_ds.get_state(**xtra)
+            state['cls'] = self.__class__
+            try_save(state, self.path, file)
+        if destroy: self.destroy()
+
+    def save(self, file:PathLikeOrBinaryStream=None, return_path:bool=False, with_opt:bool=True):
+        "Save model and optimizer state (if `with_opt`) with `file` to `self.model_dir`. `file` can be file-like (file or buffer)"
+        if is_pathlike(file): self._test_writeable_path()
+        if rank_distrib(): return # don't save if slave proc
+        target = self.path/self.model_dir/f'{file}.pth' if is_pathlike(file) else file
+        if not hasattr(self, 'opt'): with_opt=False
+        if not with_opt: state = get_model(self.model).state_dict()
+        else: state = {'model': get_model(self.model).state_dict(), 'opt':self.opt.state_dict()}
+        torch.save(state, target)
+        if return_path: return target
+
+    def dl(self, ds_type:DatasetType=DatasetType.Valid):
+        "Return DataLoader for DatasetType `ds_type`."
+        return self.data.dl(ds_type)
+
+    def load(self, file:PathLikeOrBinaryStream=None, device:torch.device=None, strict:bool=True,
+             with_opt:bool=None, purge:bool=True, remove_module:bool=False):
+        "Load model and optimizer state (if `with_opt`) `file` from `self.model_dir` using `device`. `file` can be file-like (file or buffer)"
+        if purge: self.purge(clear_opt=ifnone(with_opt, False))
+        if device is None: device = self.data.device
+        elif isinstance(device, int): device = torch.device('cuda', device)
+        source = self.path/self.model_dir/f'{file}.pth' if is_pathlike(file) else file
+        state = torch.load(source, map_location=device)
+        if set(state.keys()) == {'model', 'opt'}:
+            model_state = state['model']
+            if remove_module: model_state = remove_module_load(model_state)
+            get_model(self.model).load_state_dict(model_state, strict=strict)
+            if ifnone(with_opt,True):
+                if not hasattr(self, 'opt'): self.create_opt(defaults.lr, self.wd)
+                try:    self.opt.load_state_dict(state['opt'])
+                except: pass
+        else:
+            if with_opt: warn("Saved filed doesn't contain an optimizer state.")
+            if remove_module: state = remove_module_load(state)
+            get_model(self.model).load_state_dict(state, strict=strict)
+        del state
+        gc.collect()
+        return self
+
+    def destroy(self):
+        "Free the Learner internals, leaving just an empty shell that consumes no memory"
+
+        class ZombieLearner(Learner):
+            msg = "this object has been destroyed"
+            def __getattr__(self, item):    print(ZombieLearner.msg); return None
+            def destroyed(*args, **kwargs): print(ZombieLearner.msg)
+
+        attrs = [k for k in self.__dict__.keys() if not k.startswith("__")]
+        for a in attrs: delattr(self, a)
+        # the instance methods can still be called, but will just give a message
+        methods = [k for k in dir(self) if not k.startswith("__") and inspect.isroutine(getattr(self, k))]
+        for m in methods: setattr(self, m, ZombieLearner.destroyed)
+        self.__class__ = ZombieLearner
+        gc.collect()
+        print("this Learner object self-destroyed - it still exists, but no longer usable")
+
+    def purge(self, clear_opt:bool=True):
+        "Purge the `Learner` of all cached attributes to release some GPU memory."
+        self._test_writeable_path()
+        attrs_all = [k for k in self.__dict__.keys() if not k.startswith("__")]
+        attrs_pkl = ['bn_wd', 'callback_fns', 'layer_groups', 'loss_func', 'metrics', 'model',
+                     'model_dir', 'opt_func', 'path', 'train_bn', 'true_wd', 'wd']
+        # +callbacks: get pickled too, but not directly
+        attrs_keep = ['data', 'recorder']
+        attrs_del = list(set(attrs_all) - set(attrs_keep))
+        state = {a:getattr(self, a) for a in attrs_pkl}
+        state['cb_state'] = {cb.__class__:cb.get_state() for cb in self.callbacks}
+        if hasattr(self, 'opt'): state['opt'] = self.opt.get_state()
+
+        tmp_file = get_tmp_file(self.path/self.model_dir)
+        torch.save(state, open(tmp_file, 'wb'))
+        for a in attrs_del: delattr(self, a)
+        gc.collect()
+        state = torch.load(tmp_file)
+        os.remove(tmp_file)
+
+        for a in attrs_pkl: setattr(self, a, state[a])
+        cb_state = state.pop('cb_state')
+        self.callbacks = [load_callback(c,s, self) for c,s in cb_state.items()]
+        if not clear_opt and 'opt' in state:
+            try: self.opt = OptimWrapper.load_with_state_and_layer_group(state['opt'], self.layer_groups)
+            except: warn("Wasn't able to properly load the optimizer state again.")
+        del state
+        gc.collect()
+        return self
+
+    def get_preds(self, ds_type:DatasetType=DatasetType.Valid, with_loss:bool=False, n_batch:Optional[int]=None,
+                  pbar:Optional[PBar]=None) -> List[Tensor]:
+        "Return predictions and targets on `ds_type` dataset."
+        lf = self.loss_func if with_loss else None
+        return get_preds(self.model, self.dl(ds_type), cb_handler=CallbackHandler(self.callbacks),
+                         activ=_loss_func2activ(self.loss_func), loss_func=lf, n_batch=n_batch, pbar=pbar)
+
+    def pred_batch(self, ds_type:DatasetType=DatasetType.Valid, batch:Tuple=None, reconstruct:bool=False, with_dropout:bool=False) -> List[Tensor]:
+        with torch.no_grad():
+            training = self.model.training
+            self.model.train(False)
+            "Return output of the model on one batch from `ds_type` dataset."
+            if batch is not None: xb,yb = batch
+            else: xb,yb = self.data.one_batch(ds_type, detach=False, denorm=False)
+            cb_handler = CallbackHandler(self.callbacks)
+            xb,yb = cb_handler.on_batch_begin(xb,yb, train=False)
+            if not with_dropout: 
+                preds = loss_batch(self.model.eval(), xb, yb, cb_handler=cb_handler)
+            else: 
+                preds = loss_batch(self.model.eval().apply(self.apply_dropout), xb, yb, cb_handler=cb_handler)
+            res = _loss_func2activ(self.loss_func)(preds[0])
+            self.model.train(training)
+            if not reconstruct: return res
+            res = res.detach().cpu()
+            ds = self.dl(ds_type).dataset
+            norm = getattr(self.data, 'norm', False)
+            if norm and norm.keywords.get('do_y',False):
+                res = self.data.denorm(res, do_x=True)
+            return [ds.reconstruct(o) for o in res]
+
+    def backward(self, item):
+        "Pass `item` through the model and computes the gradient. Useful if `backward_hooks` are attached."
+        xb,yb = self.data.one_item(item)
+        loss = loss_batch(self.model.eval(), xb, yb, self.loss_func, opt=FakeOptimizer(),
+                          cb_handler=CallbackHandler(self.callbacks))
+        return loss
+
+    def predict(self, item:ItemBase, return_x:bool=False, batch_first:bool=True, with_dropout:bool=False, **kwargs):
+        "Return predicted class, label and probabilities for `item`."
+        batch = self.data.one_item(item)
+        res = self.pred_batch(batch=batch, with_dropout=with_dropout)
+        raw_pred,x = grab_idx(res,0,batch_first=batch_first),batch[0]
+        norm = getattr(self.data,'norm',False)
+        if norm:
+            x = self.data.denorm(x)
+            if norm.keywords.get('do_y',False): raw_pred = self.data.denorm(raw_pred)
+        ds = self.data.single_ds
+        pred = ds.y.analyze_pred(raw_pred, **kwargs)
+        x = ds.x.reconstruct(grab_idx(x, 0))
+        y = ds.y.reconstruct(pred, x) if has_arg(ds.y.reconstruct, 'x') else ds.y.reconstruct(pred)
+        return (x, y, pred, raw_pred) if return_x else (y, pred, raw_pred)
+
+    def validate(self, dl=None, callbacks=None, metrics=None):
+        "Validate on `dl` with potential `callbacks` and `metrics`."
+        dl = ifnone(dl, self.data.valid_dl)
+        metrics = ifnone(metrics, self.metrics)
+        cb_handler = CallbackHandler(self.callbacks + ifnone(callbacks, []), metrics)
+        cb_handler.on_epoch_begin()
+        val_metrics = validate(self.model, dl, self.loss_func, cb_handler)
+        cb_handler.on_epoch_end(val_metrics)
+        return cb_handler.state_dict['last_metrics']
+
+    def show_results(self, ds_type=DatasetType.Valid, rows:int=5, **kwargs):
+        "Show `rows` result of predictions on `ds_type` dataset."
+        #TODO: get read of has_arg x and split_kwargs_by_func if possible
+        #TODO: simplify this and refactor with pred_batch(...reconstruct=True)
+        n_items = rows ** 2 if self.data.train_ds.x._square_show_res else rows
+        if self.dl(ds_type).batch_size < n_items: n_items = self.dl(ds_type).batch_size
+        ds = self.dl(ds_type).dataset
+        self.callbacks.append(RecordOnCPU())
+        preds = self.pred_batch(ds_type)
+        *self.callbacks,rec_cpu = self.callbacks
+        x,y = rec_cpu.input,rec_cpu.target
+        norm = getattr(self.data,'norm',False)
+        if norm:
+            x = self.data.denorm(x)
+            if norm.keywords.get('do_y',False):
+                y     = self.data.denorm(y, do_x=True)
+                preds = self.data.denorm(preds, do_x=True)
+        analyze_kwargs,kwargs = split_kwargs_by_func(kwargs, ds.y.analyze_pred)
+        preds = [ds.y.analyze_pred(grab_idx(preds, i), **analyze_kwargs) for i in range(n_items)]
+        xs = [ds.x.reconstruct(grab_idx(x, i)) for i in range(n_items)]
+        if has_arg(ds.y.reconstruct, 'x'):
+            ys = [ds.y.reconstruct(grab_idx(y, i), x=x) for i,x in enumerate(xs)]
+            zs = [ds.y.reconstruct(z, x=x) for z,x in zip(preds,xs)]
+        else :
+            ys = [ds.y.reconstruct(grab_idx(y, i)) for i in range(n_items)]
+            zs = [ds.y.reconstruct(z) for z in preds]
+        ds.x.show_xyzs(xs, ys, zs, **kwargs)
+
+    def apply_dropout(self, m):
+        "If a module contains 'dropout' in it's name, it will be switched to .train() mode."
+        if 'dropout' in m.__class__.__name__.lower(): m.train()
+
+    def predict_with_mc_dropout(self, item:ItemBase, with_dropout:bool=True, n_times=10, **kwargs):
+        "Make predictions with dropout turned on for n_times (default 10)."
+        return [self.predict(item, with_dropout=with_dropout) for _ in range(n_times)]
+
+class RecordOnCPU(Callback):
+    "Store the `input` and `target` going through the model on the CPU."
+    def on_batch_begin(self, last_input,last_target,**kwargs):
+        self.input,self.target = to_cpu(last_input),to_cpu(last_target)
+
+class LearnerCallback(Callback):
+    "Base class for creating callbacks for a `Learner`."
+    def __init__(self, learn):
+        self._learn = weakref.ref(learn)
+        self.exclude,self.not_min = ['_learn'],[]
+        setattr(self.learn, self.cb_name, self)
+
+    def __getattr__(self,k): return getattr(self.learn, k)
+    def __setstate__(self,data:Any): self.__dict__.update(data)
+
+    @property
+    def learn(self) -> Learner: return self._learn()
+    @learn.setter
+    def learn(self, learn: Learner) -> None: self._learn = weakref.ref(learn)
+
+    @property
+    def cb_name(self): return camel2snake(self.__class__.__name__)
+
+class Recorder(LearnerCallback):
+    "A `LearnerCallback` that records epoch, loss, opt and metric data during training."
+    _order=-10
+    def __init__(self, learn:Learner, add_time:bool=True, silent:bool=False):
+        super().__init__(learn)
+        self.opt = self.learn.opt
+        self.train_dl = self.learn.data.train_dl
+        self.no_val,self.silent,self.add_time = False,silent,add_time
+
+    def on_train_begin(self, pbar:PBar, metrics_names:Collection[str], **kwargs:Any)->None:
+        "Initialize recording status at beginning of training."
+        self.pbar = pbar
+        self.names = ['epoch', 'train_loss'] if self.no_val else ['epoch', 'train_loss', 'valid_loss']
+        self.metrics_names = metrics_names
+        if hasattr(self, '_added_met_names'): self.metrics_names += self._added_met_names
+        self.names += self.metrics_names
+        if self.add_time: self.names.append('time')
+        if not self.silent: self.pbar.write(self.names, table=True)
+        self.losses,self.val_losses,self.lrs,self.moms,self.metrics,self.nb_batches = [],[],[],[],[],[]
+
+    def on_epoch_begin(self, **kwargs:Any)->None:
+        if self.add_time: self.start_epoch = time()
+
+    def on_batch_begin(self, train, **kwargs:Any)->None:
+        "Record learning rate and momentum at beginning of batch."
+        if train:
+            self.lrs.append(self.opt.lr)
+            self.moms.append(self.opt.mom)
+
+    def on_backward_begin(self, smooth_loss:Tensor, **kwargs:Any)->None:
+        "Record the loss before any other callback has a chance to modify it."
+        self.losses.append(smooth_loss)
+        if self.pbar is not None and hasattr(self.pbar,'child'):
+            self.pbar.child.comment = f'{smooth_loss:.4f}'
+
+    def on_epoch_end(self, epoch:int, num_batch:int, smooth_loss:Tensor,
+                     last_metrics=MetricsList, **kwargs:Any)->bool:
+        "Save epoch info: num_batch, smooth_loss, metrics."
+        self.nb_batches.append(num_batch)
+        if last_metrics is not None: self.val_losses.append(last_metrics[0])
+        else: last_metrics = [] if self.no_val else [None]
+        if len(last_metrics) > 1: self.metrics.append(last_metrics[1:])
+        self.format_stats([epoch, smooth_loss] + last_metrics)
+
+    def format_stats(self, stats:TensorOrNumList)->None:
+        "Format stats before printing."
+        str_stats = []
+        for name,stat in zip(self.names,stats):
+            str_stats.append('#na#' if stat is None else str(stat) if isinstance(stat, int) else f'{stat:.6f}')
+        if self.add_time: str_stats.append(format_time(time() - self.start_epoch))
+        if not self.silent: self.pbar.write(str_stats, table=True)
+
+    def add_metric_names(self, names):
+        "Add `names` to the inner metric names."
+        if hasattr(self, '_added_met_names'): self._added_met_names += names
+        else:                                 self._added_met_names  = names
+
+    def plot_lr(self, show_moms=False, skip_start:int=0, skip_end:int=0, return_fig:bool=None)->Optional[plt.Figure]:
+        "Plot learning rate, `show_moms` to include momentum."
+        lrs = self._split_list(self.lrs, skip_start, skip_end)
+        iterations = self._split_list(range_of(self.lrs), skip_start, skip_end)
+        if show_moms:
+            moms = self._split_list(self.moms, skip_start, skip_end)
+            fig, axs = plt.subplots(1,2, figsize=(12,4))
+            axs[0].plot(iterations, lrs)
+            axs[0].set_xlabel('Iterations')
+            axs[0].set_ylabel('Learning Rate')
+            axs[1].plot(iterations, moms)
+            axs[1].set_xlabel('Iterations')
+            axs[1].set_ylabel('Momentum')
+        else:
+            fig, ax = plt.subplots()
+            ax.plot(iterations, lrs)
+            ax.set_xlabel('Iterations')
+            ax.set_ylabel('Learning Rate')
+        if ifnone(return_fig, defaults.return_fig): return fig
+        if not IN_NOTEBOOK: plot_sixel(fig)
+
+    @staticmethod
+    def smoothen_by_spline(xs, ys, **kwargs):
+        xs = np.arange(len(ys))
+        spl = scipy.interpolate.UnivariateSpline(xs, ys, **kwargs)
+        ys = spl(xs)
+        return ys
+
+    def plot(self, skip_start:int=10, skip_end:int=5, suggestion:bool=False, return_fig:bool=None,
+             **kwargs)->Optional[plt.Figure]:
+        "Plot learning rate and losses, trimmed between `skip_start` and `skip_end`. Optionally plot and return min gradient"
+        lrs = self._split_list(self.lrs, skip_start, skip_end)
+        losses = self._split_list(self.losses, skip_start, skip_end)
+        losses = [x.item() for x in losses]
+        if 'k' in kwargs: losses = self.smoothen_by_spline(lrs, losses, **kwargs)
+        fig, ax = plt.subplots(1,1)
+        ax.plot(lrs, losses)
+        ax.set_ylabel("Loss")
+        ax.set_xlabel("Learning Rate")
+        ax.set_xscale('log')
+        ax.xaxis.set_major_formatter(plt.FormatStrFormatter('%.0e'))
+        if suggestion:
+            try: mg = (np.gradient(np.array(losses))).argmin()
+            except:
+                print("Failed to compute the gradients, there might not be enough points.")
+                return
+            print(f"Min numerical gradient: {lrs[mg]:.2E}")
+            ax.plot(lrs[mg],losses[mg],markersize=10,marker='o',color='red')
+            self.min_grad_lr = lrs[mg]
+            ml = np.argmin(losses)
+            print(f"Min loss divided by 10: {lrs[ml]/10:.2E}")
+        if ifnone(return_fig, defaults.return_fig): return fig
+        if not IN_NOTEBOOK: plot_sixel(fig)
+
+    def plot_losses(self, skip_start:int=0, skip_end:int=0, return_fig:bool=None)->Optional[plt.Figure]:
+        "Plot training and validation losses."
+        fig, ax = plt.subplots(1,1)
+        losses = self._split_list(self.losses, skip_start, skip_end)
+        iterations = self._split_list(range_of(self.losses), skip_start, skip_end)
+        ax.plot(iterations, losses, label='Train')
+        val_iter = self._split_list_val(np.cumsum(self.nb_batches), skip_start, skip_end)
+        val_losses = self._split_list_val(self.val_losses, skip_start, skip_end)
+        ax.plot(val_iter, val_losses, label='Validation')
+        ax.set_ylabel('Loss')
+        ax.set_xlabel('Batches processed')
+        ax.legend()
+        if ifnone(return_fig, defaults.return_fig): return fig
+        if not IN_NOTEBOOK: plot_sixel(fig)
+
+    def plot_metrics(self, skip_start:int=0, skip_end:int=0, return_fig:bool=None)->Optional[plt.Figure]:
+        "Plot metrics collected during training."
+        assert len(self.metrics) != 0, "There are no metrics to plot."
+        fig, axes = plt.subplots(len(self.metrics[0]),1,figsize=(6, 4*len(self.metrics[0])))
+        val_iter = self._split_list_val(np.cumsum(self.nb_batches), skip_start, skip_end)
+        axes = axes.flatten() if len(self.metrics[0]) != 1 else [axes]
+        for i, ax in enumerate(axes):
+            values = [met[i] for met in self.metrics]
+            values = self._split_list_val(values, skip_start, skip_end)
+            ax.plot(val_iter, values)
+            ax.set_ylabel(str(self.metrics_names[i]))
+            ax.set_xlabel('Batches processed')             
+        if ifnone(return_fig, defaults.return_fig): return fig
+        if not IN_NOTEBOOK: plot_sixel(fig)
+
+    def _split_list(self, vals:Collection[float], skip_start:int, skip_end:int):
+        return vals[skip_start:-skip_end] if skip_end > 0 else vals[skip_start:]
+
+    def _split_list_val(self, vals:Collection[float], skip_start:int, skip_end:int):
+        val_iter = np.cumsum(self.nb_batches)
+        start_val = (val_iter - skip_start >= 0).nonzero()[0].min()
+        end_val = (val_iter[-1] - val_iter - skip_end >= 0).nonzero()[0].max()+1
+        return vals[start_val:end_val] if skip_end > 0 else vals[start_val:]
+
+class FakeOptimizer():
+    def step(self): pass
+    def zero_grad(self): pass
+
+def load_callback(class_func, state, learn:Learner):
+    init_kwargs, others = split_kwargs_by_func(state, class_func.__init__)
+    res = class_func(learn, **init_kwargs) if issubclass(class_func, LearnerCallback) else class_func(**init_kwargs)
+    for k,v in others.items(): setattr(res, k, v)
+    return res
+
+def load_learner(path:PathOrStr, file:PathLikeOrBinaryStream='export.pkl', test:ItemList=None, **db_kwargs):
+    "Load a `Learner` object saved with `export_state` in `path/file` with empty data, optionally add `test` and load on `cpu`. `file` can be file-like (file or buffer)"
+    source = Path(path)/file if is_pathlike(file) else file
+    state = torch.load(source, map_location='cpu') if defaults.device == torch.device('cpu') else torch.load(source)
+    model = state.pop('model')
+    src = LabelLists.load_state(path, state.pop('data'))
+    if test is not None: src.add_test(test)
+    data = src.databunch(**db_kwargs)
+    cb_state = state.pop('cb_state')
+    clas_func = state.pop('cls')
+    res = clas_func(data, model, **state)
+    res.callback_fns = state['callback_fns'] #to avoid duplicates
+    res.callbacks = [load_callback(c,s, res) for c,s in cb_state.items()]
+    return res

DeOldify/fastai/basics.py

@@ -0,0 +1,29 @@
+from .basic_train import *
+from .callback import *
+from .core import *
+from .basic_data import *
+from .data_block import *
+from .layers import *
+from .metrics import *
+from .torch_core import *
+from .train import *
+from .datasets import *
+from .version import *
+from . import callbacks
+
+"""
+from . import core,torch_core,basic_data,basic_train,callback,data_block,layers,metrics,train,datasets,callbacks
+
+__all__  = [o for o in dir(core) if not o.startswith('_')]
+__all__ += [o for o in dir(torch_core) if not o.startswith('_')]
+__all__ += [*basic_train.__all__, *callback.__all__, 'core', 'torch_core', 'callbacks',
+           *basic_data.__all__, *data_block.__all__, *layers.__all__, *metrics.__all__,
+           *train.__all__, *datasets.__all__, '__version__']
+"""
+
+try: from .gen_doc.nbdoc import doc
+except: pass  # Optional if jupyter is present
+    #__all__.append('doc')
+
+__all__ = [o for o in dir(sys.modules[__name__]) if not o.startswith('_')] + ['__version__']
+

DeOldify/fastai/callback.py

@@ -0,0 +1,396 @@
+"Callbacks provides extensibility to the `basic_train` loop. See `train` for examples of custom callbacks."
+from .basic_data import *
+from .torch_core import *
+import torch.distributed as dist
+
+__all__ = ['AverageMetric', 'Callback', 'CallbackHandler', 'OptimWrapper', 'SmoothenValue', 'Scheduler', 'annealing_cos', 'CallbackList',
+           'annealing_exp', 'annealing_linear', 'annealing_no', 'annealing_poly']
+
+class OptimWrapper():
+    "Basic wrapper around `opt` to simplify hyper-parameters changes."
+    def __init__(self, opt:optim.Optimizer, wd:Floats=0., true_wd:bool=False, bn_wd:bool=True):
+        assert not isinstance(opt, OptimWrapper)
+        self.opt,self.true_wd,self.bn_wd = opt,true_wd,bn_wd
+        self.opt_keys = list(self.opt.param_groups[0].keys())
+        self.opt_keys.remove('params')
+        self.read_defaults()
+        self.wd = wd
+
+    @classmethod
+    def create(cls, opt_func:Union[type,Callable], lr:Union[float,Tuple,List], layer_groups:ModuleList, wd:Floats=0., 
+               true_wd:bool=False, bn_wd:bool=True)->optim.Optimizer:
+        "Create an `optim.Optimizer` from `opt_func` with `lr`. Set lr on `layer_groups`."
+        split_params = split_no_wd_params(layer_groups)
+        opt = opt_func([{'params': p, 'lr':0} for p in split_params])
+        opt = cls(opt, wd=wd, true_wd=true_wd, bn_wd=bn_wd)
+        opt.lr,opt.opt_func = listify(lr, layer_groups),opt_func
+        return opt
+
+    def new(self, layer_groups:Collection[nn.Module], split_no_wd:bool=True):
+        "Create a new `OptimWrapper` from `self` with another `layer_groups` but the same hyper-parameters."
+        opt_func = getattr(self, 'opt_func', self.opt.__class__)
+        res = self.create(opt_func, self.lr, layer_groups, wd=self.wd, true_wd=self.true_wd, bn_wd=self.bn_wd)
+        res.mom,res.beta = self.mom,self.beta
+        return res
+
+    def new_with_params(self, param_groups:Collection[Collection[nn.Parameter]]):
+        "Create a new `OptimWrapper` from `self` with another `layer_groups` but the same hyper-parameters."
+        opt_func = getattr(self, 'opt_func', self.opt.__class__)
+        opt = opt_func([{'params': p, 'lr':0} for p in param_groups])
+        opt = self.__class__(opt, wd=self.wd, true_wd=self.true_wd, bn_wd=self.bn_wd)
+        opt.lr,opt.opt_func,opt.mom,opt.beta = self.lr,opt_func,self.mom,self.beta
+        return opt
+
+    def __repr__(self)->str:
+        return f'OptimWrapper over {repr(self.opt)}.\nTrue weight decay: {self.true_wd}'
+
+    #Pytorch optimizer methods
+    def step(self)->None:
+        "Set weight decay and step optimizer."
+        # weight decay outside of optimizer step (AdamW)
+        if self.true_wd:
+            for lr,wd,pg1,pg2 in zip(self._lr,self._wd,self.opt.param_groups[::2],self.opt.param_groups[1::2]):
+                for p in pg1['params']: p.data.mul_(1 - wd*lr)
+                if self.bn_wd:
+                    for p in pg2['params']: p.data.mul_(1 - wd*lr)
+            self.set_val('weight_decay', listify(0, self._wd))
+        self.opt.step()
+
+    def zero_grad(self)->None:
+        "Clear optimizer gradients."
+        self.opt.zero_grad()
+
+    #Passthrough to the inner opt.
+    def __getattr__(self, k:str)->Any: return getattr(self.opt, k, None)
+    def __setstate__(self,data:Any): self.__dict__.update(data)
+
+    def clear(self):
+        "Reset the state of the inner optimizer."
+        sd = self.state_dict()
+        sd['state'] = {}
+        self.load_state_dict(sd)
+
+    @property
+    def n_params(self): return sum([len(pg['params']) for pg in self.opt.param_groups])
+
+    #Hyperparameters as properties
+    @property
+    def lr(self)->float: return self._lr[-1]
+    @lr.setter
+    def lr(self, val:float)->None:
+        self._lr = self.set_val('lr', listify(val, self._lr))
+
+    @property
+    def mom(self)->float:return self._mom[-1]
+    @mom.setter
+    def mom(self, val:float)->None:
+        if 'momentum' in self.opt_keys: self.set_val('momentum', listify(val, self._mom))
+        elif 'betas' in self.opt_keys:  self.set_val('betas', (listify(val, self._mom), self._beta))
+        self._mom = listify(val, self._mom)
+
+    @property
+    def beta(self)->float: return None if self._beta is None else self._beta[-1]
+    @beta.setter
+    def beta(self, val:float)->None:
+        "Set beta (or alpha as makes sense for given optimizer)."
+        if val is None: return
+        if 'betas' in self.opt_keys:    self.set_val('betas', (self._mom, listify(val, self._beta)))
+        elif 'alpha' in self.opt_keys:  self.set_val('alpha', listify(val, self._beta))
+        self._beta = listify(val, self._beta)
+
+    @property
+    def wd(self)->float: return self._wd[-1]
+    @wd.setter
+    def wd(self, val:float)->None:
+        "Set weight decay."
+        if not self.true_wd: self.set_val('weight_decay', listify(val, self._wd), bn_groups=self.bn_wd)
+        self._wd = listify(val, self._wd)
+
+    #Helper functions
+    def read_defaults(self)->None:
+        "Read the values inside the optimizer for the hyper-parameters."
+        self._beta = None
+        if 'lr' in self.opt_keys: self._lr = self.read_val('lr')
+        if 'momentum' in self.opt_keys: self._mom = self.read_val('momentum')
+        if 'alpha' in self.opt_keys: self._beta = self.read_val('alpha')
+        if 'betas' in self.opt_keys: self._mom,self._beta = self.read_val('betas')
+        if 'weight_decay' in self.opt_keys: self._wd = self.read_val('weight_decay')
+        reserved_names = ['params', 'lr', 'momentum', 'alpha', 'betas', 'weight_decay']
+        stat_names = [n for n in self.opt_keys if n not in reserved_names]
+        self._stats = {n:self.read_val(n) for n in stat_names}
+
+    def get_stat(self, name:str)->float: 
+        if name in ['lr', 'mom', 'beta', 'wd']: return getattr(self, name)
+        else: return self._stats[name][-1]
+    def set_stat(self, name:str, value:Union[float, Collection[float]])->None:
+        if name in ['lr', 'mom', 'beta', 'wd']: setattr(self, name, value)
+        else:
+            val = listify(value, self._stats[name])
+            self.set_val(name, val)
+            self._stats[name] = val
+
+    def set_val(self, key:str, val:Any, bn_groups:bool=True)->Any:
+        "Set `val` inside the optimizer dictionary at `key`."
+        if is_tuple(val): val = [(v1,v2) for v1,v2 in zip(*val)]
+        for v,pg1,pg2 in zip(val,self.opt.param_groups[::2],self.opt.param_groups[1::2]):
+            pg1[key] = v
+            if bn_groups: pg2[key] = v
+        return val
+
+    def read_val(self, key:str) -> Union[List[float],Tuple[List[float],List[float]]]:
+        "Read a hyperparameter `key` in the optimizer dictionary."
+        val = [pg[key] for pg in self.opt.param_groups[::2]]
+        if is_tuple(val[0]): val = [o[0] for o in val], [o[1] for o in val]
+        return val
+    
+    def get_state(self):
+        "Return the inner state minus the layer groups."
+        return {'opt_state':self.opt.state_dict(), 'lr':self._lr, 'wd':self._wd, 'beta':self._beta, 'mom':self._mom,
+                'opt_func':self.opt_func, 'true_wd':self.true_wd, 'bn_wd':self.bn_wd}
+
+    @classmethod
+    def load_with_state_and_layer_group(cls, state:dict, layer_groups:Collection[nn.Module]):
+        res = cls.create(state['opt_func'], state['lr'], layer_groups, wd=state['wd'], true_wd=state['true_wd'], 
+                     bn_wd=state['bn_wd'])
+        res._mom,res._beta = state['mom'],state['beta']
+        res.load_state_dict(state['opt_state'])
+        return res
+
+class Callback():
+    "Base class for callbacks that want to record values, dynamically change learner params, etc."
+    _order=0
+    def on_train_begin(self, **kwargs:Any)->None:
+        "To initialize constants in the callback."
+        pass
+    def on_epoch_begin(self, **kwargs:Any)->None:
+        "At the beginning of each epoch."
+        pass
+    def on_batch_begin(self, **kwargs:Any)->None:
+        "Set HP before the output and loss are computed."
+        pass
+    def on_loss_begin(self, **kwargs:Any)->None:
+        "Called after forward pass but before loss has been computed."
+        pass
+    def on_backward_begin(self, **kwargs:Any)->None:
+        "Called after the forward pass and the loss has been computed, but before backprop."
+        pass
+    def on_backward_end(self, **kwargs:Any)->None:
+        "Called after backprop but before optimizer step. Useful for true weight decay in AdamW."
+        pass
+    def on_step_end(self, **kwargs:Any)->None:
+        "Called after the step of the optimizer but before the gradients are zeroed."
+        pass
+    def on_batch_end(self, **kwargs:Any)->None:
+        "Called at the end of the batch."
+        pass
+    def on_epoch_end(self, **kwargs:Any)->None:
+        "Called at the end of an epoch."
+        pass
+    def on_train_end(self, **kwargs:Any)->None:
+        "Useful for cleaning up things and saving files/models."
+        pass
+    def jump_to_epoch(self, epoch)->None:
+        "To resume training at `epoch` directly."
+        pass
+
+    def get_state(self, minimal:bool=True):
+        "Return the inner state of the `Callback`, `minimal` or not."
+        to_remove = ['exclude', 'not_min'] + getattr(self, 'exclude', []).copy()
+        if minimal: to_remove += getattr(self, 'not_min', []).copy()
+        return {k:v for k,v in self.__dict__.items() if k not in to_remove}
+
+    def  __repr__(self):
+        attrs = func_args(self.__init__)
+        to_remove = getattr(self, 'exclude', [])
+        list_repr = [self.__class__.__name__] + [f'{k}: {getattr(self, k)}' for k in attrs if k != 'self' and k not in to_remove]
+        return '\n'.join(list_repr)
+
+class SmoothenValue():
+    "Create a smooth moving average for a value (loss, etc) using `beta`."
+    def __init__(self, beta:float):
+        self.beta,self.n,self.mov_avg = beta,0,0
+
+    def add_value(self, val:float)->None:
+        "Add `val` to calculate updated smoothed value."
+        self.n += 1
+        self.mov_avg = self.beta * self.mov_avg + (1 - self.beta) * val
+        self.smooth = self.mov_avg / (1 - self.beta ** self.n)
+
+CallbackList = Collection[Callback]
+
+def _get_init_state(): return {'epoch':0, 'iteration':0, 'num_batch':0, 'skip_validate': False}
+
+@dataclass
+class CallbackHandler():
+    "Manage all of the registered `callbacks` and `metrics`, smoothing loss by momentum `beta`."
+    callbacks:CallbackList=None
+    metrics:CallbackList=None
+    beta:float=0.98
+
+    def __post_init__(self)->None:
+        "Initialize smoother and learning stats."
+        self.callbacks = ifnone(self.callbacks, [])
+        self.metrics = ifnone(self.metrics, [])
+        self.metrics = [(met if isinstance(met, Callback) else AverageMetric(met)) for met in self.metrics]
+        self.callbacks = sorted(self.callbacks, key=lambda o: getattr(o, '_order', 0))
+        self.smoothener = SmoothenValue(self.beta)
+        self.state_dict:Dict[str,Union[int,float,Tensor]]=_get_init_state()
+
+    def _call_and_update(self, cb, cb_name, **kwargs)->None:
+        "Call `cb_name` on `cb` and update the inner state."
+        new = ifnone(getattr(cb, f'on_{cb_name}')(**self.state_dict, **kwargs), dict())
+        for k,v in new.items():
+            if k not in self.state_dict:
+                raise Exception(f"{k} isn't a valid key in the state of the callbacks.")
+            else: self.state_dict[k] = v
+    
+    def __call__(self, cb_name, call_mets=True, **kwargs)->None:
+        "Call through to all of the `CallbakHandler` functions."
+        if call_mets: 
+            for met in self.metrics: self._call_and_update(met, cb_name, **kwargs)
+        for cb in self.callbacks: self._call_and_update(cb, cb_name, **kwargs)
+
+    def set_dl(self, dl:DataLoader):
+        "Set the current `dl` used."
+        if hasattr(self, 'cb_dl'): self.callbacks.remove(self.cb_dl)
+        if isinstance(dl.dataset, Callback):
+            self.callbacks.append(dl.dataset)
+            self.cb_dl = dl.dataset
+
+    def on_train_begin(self, epochs:int, pbar:PBar, metrics:MetricFuncList)->None:
+        "About to start learning."
+        self.state_dict = _get_init_state()
+        self.state_dict.update(dict(n_epochs=epochs, pbar=pbar, metrics=metrics))
+        names = [(met.name if hasattr(met, 'name') else camel2snake(met.__class__.__name__)) for met in self.metrics]
+        self('train_begin', metrics_names=names)
+        if self.state_dict['epoch'] != 0:
+            self.state_dict['pbar'].first_bar.total -= self.state_dict['epoch']
+            for cb in self.callbacks: cb.jump_to_epoch(self.state_dict['epoch'])
+
+    def on_epoch_begin(self)->None:
+        "Handle new epoch."
+        self.state_dict['num_batch'],self.state_dict['stop_training'] = 0,False
+        self('epoch_begin')
+
+    def on_batch_begin(self, xb:Tensor, yb:Tensor, train:bool=True)->Tuple[Any,Any]:
+        "Handle new batch `xb`,`yb` in `train` or validation."
+        self.state_dict.update(dict(last_input=xb, last_target=yb, train=train, 
+            stop_epoch=False, skip_step=False, skip_zero=False, skip_bwd=False))
+        self('batch_begin', mets = not self.state_dict['train'])
+        return self.state_dict['last_input'], self.state_dict['last_target']
+
+    def on_loss_begin(self, out:Tensor)->Any:
+        "Handle start of loss calculation with model output `out`."
+        self.state_dict['last_output'] = out
+        self('loss_begin', call_mets=False)
+        return self.state_dict['last_output']
+
+    def on_backward_begin(self, loss:Tensor)->Tuple[Any,Any]:
+        "Handle gradient calculation on `loss`."
+        self.smoothener.add_value(loss.detach().cpu())
+        self.state_dict['last_loss'], self.state_dict['smooth_loss'] = loss, self.smoothener.smooth
+        self('backward_begin', call_mets=False)
+        return self.state_dict['last_loss'], self.state_dict['skip_bwd']
+
+    def on_backward_end(self)->Any:
+        "Handle end of gradient calculation."
+        self('backward_end', call_mets=False)
+        return self.state_dict['skip_step']
+
+    def on_step_end(self)->Any:
+        "Handle end of optimization step."
+        self('step_end', call_mets=False)
+        return self.state_dict['skip_zero']
+
+    def on_batch_end(self, loss:Tensor)->Any:
+        "Handle end of processing one batch with `loss`."
+        self.state_dict['last_loss'] = loss
+        self('batch_end', call_mets = not self.state_dict['train'])
+        if self.state_dict['train']:
+            self.state_dict['iteration'] += 1
+            self.state_dict['num_batch'] += 1
+        return self.state_dict['stop_epoch']
+
+    def on_epoch_end(self, val_loss:Tensor)->bool:
+        "Epoch is done, process `val_loss`."
+        self.state_dict['last_metrics'] = [val_loss] if val_loss is not None else [None]
+        self('epoch_end', call_mets = val_loss is not None)
+        self.state_dict['epoch'] += 1
+        return self.state_dict['stop_training']
+
+    def on_train_end(self, exception:Union[bool,Exception])->None:
+        "Handle end of training, `exception` is an `Exception` or False if no exceptions during training."
+        self('train_end', exception=exception)
+        
+    @property
+    def skip_validate(self): return self.state_dict['skip_validate']
+
+class AverageMetric(Callback):
+    "Wrap a `func` in a callback for metrics computation."
+    def __init__(self, func):
+        # If func has a __name__ use this one else it should be a partial
+        name = func.__name__ if hasattr(func, '__name__') else func.func.__name__
+        self.func, self.name = func, name
+        self.world = num_distrib()
+
+    def on_epoch_begin(self, **kwargs):
+        "Set the inner value to 0."
+        self.val, self.count = 0.,0
+
+    def on_batch_end(self, last_output, last_target, **kwargs):
+        "Update metric computation with `last_output` and `last_target`."
+        if not is_listy(last_target): last_target=[last_target]
+        self.count += first_el(last_target).size(0)
+        val = self.func(last_output, *last_target)
+        if self.world:
+            val = val.clone()
+            dist.all_reduce(val, op=dist.ReduceOp.SUM)
+            val /= self.world
+        self.val += first_el(last_target).size(0) * val.detach().cpu()
+
+    def on_epoch_end(self, last_metrics, **kwargs):
+        "Set the final result in `last_metrics`."
+        return add_metrics(last_metrics, self.val/self.count)
+
+def annealing_no(start:Number, end:Number, pct:float)->Number:
+    "No annealing, always return `start`."
+    return start
+def annealing_linear(start:Number, end:Number, pct:float)->Number:
+    "Linearly anneal from `start` to `end` as pct goes from 0.0 to 1.0."
+    return start + pct * (end-start)
+def annealing_exp(start:Number, end:Number, pct:float)->Number:
+    "Exponentially anneal from `start` to `end` as pct goes from 0.0 to 1.0."
+    return start * (end/start) ** pct
+def annealing_cos(start:Number, end:Number, pct:float)->Number:
+    "Cosine anneal from `start` to `end` as pct goes from 0.0 to 1.0."
+    cos_out = np.cos(np.pi * pct) + 1
+    return end + (start-end)/2 * cos_out
+
+def do_annealing_poly(start:Number, end:Number, pct:float, degree:Number)->Number:
+    "Helper function for `anneal_poly`."
+    return end + (start-end) * (1-pct)**degree
+def annealing_poly(degree:Number)->Number:
+    "Anneal polynomically from `start` to `end` as pct goes from 0.0 to 1.0."
+    return functools.partial(do_annealing_poly, degree=degree)
+
+class Scheduler():
+    "Used to \"step\" from start,end (`vals`) over `n_iter` iterations on a schedule defined by `func`"
+    def __init__(self, vals:StartOptEnd, n_iter:int, func:Optional[AnnealFunc]=None):
+        self.start,self.end = (vals[0],vals[1]) if is_tuple(vals) else (vals,0)
+        self.n_iter = max(1,n_iter)
+        if func is None: self.func = annealing_linear if is_tuple(vals) else annealing_no
+        else:          self.func = func
+        self.n = 0
+        
+    def restart(self): self.n = 0
+
+    def step(self)->Number:
+        "Return next value along annealed schedule."
+        self.n += 1
+        return self.func(self.start, self.end, self.n/self.n_iter)
+
+    @property
+    def is_done(self)->bool:
+        "Return `True` if schedule completed."
+        return self.n >= self.n_iter
+

DeOldify/fastai/callbacks/__init__.py

@@ -0,0 +1,11 @@
+from .lr_finder import *
+from .one_cycle import *
+from .fp16 import *
+from .general_sched import *
+from .hooks import *
+from .mixup import *
+from .rnn import *
+from .tracker import *
+from .csv_logger import *
+from .loss_metrics import *
+from .oversampling import *

DeOldify/fastai/callbacks/csv_logger.py

@@ -0,0 +1,43 @@
+"A `Callback` that saves tracked metrics into a persistent file."
+#Contribution from devforfu: https://nbviewer.jupyter.org/gist/devforfu/ea0b3fcfe194dad323c3762492b05cae
+from ..torch_core import *
+from ..basic_data import DataBunch
+from ..callback import *
+from ..basic_train import Learner, LearnerCallback
+from time import time
+from fastprogress.fastprogress import format_time
+
+__all__ = ['CSVLogger']
+
+class CSVLogger(LearnerCallback):
+    "A `LearnerCallback` that saves history of metrics while training `learn` into CSV `filename`."
+    def __init__(self, learn:Learner, filename: str = 'history', append: bool = False): 
+        super().__init__(learn)
+        self.filename,self.path,self.append = filename,self.learn.path/f'{filename}.csv',append
+        self.add_time = True
+
+    def read_logged_file(self):  
+        "Read the content of saved file"
+        return pd.read_csv(self.path)
+
+    def on_train_begin(self, **kwargs: Any) -> None:
+        "Prepare file with metric names."
+        self.path.parent.mkdir(parents=True, exist_ok=True)      
+        self.file = self.path.open('a') if self.append else self.path.open('w')
+        self.file.write(','.join(self.learn.recorder.names[:(None if self.add_time else -1)]) + '\n')
+    
+    def on_epoch_begin(self, **kwargs:Any)->None:
+        if self.add_time: self.start_epoch = time()
+        
+    def on_epoch_end(self, epoch: int, smooth_loss: Tensor, last_metrics: MetricsList, **kwargs: Any) -> bool:
+        "Add a line with `epoch` number, `smooth_loss` and `last_metrics`."
+        last_metrics = ifnone(last_metrics, [])
+        stats = [str(stat) if isinstance(stat, int) else '#na#' if stat is None else f'{stat:.6f}'
+                 for name, stat in zip(self.learn.recorder.names, [epoch, smooth_loss] + last_metrics)]
+        if self.add_time: stats.append(format_time(time() - self.start_epoch))
+        str_stats = ','.join(stats)
+        self.file.write(str_stats + '\n')
+
+    def on_train_end(self, **kwargs: Any) -> None:  
+        "Close the file."
+        self.file.close()

DeOldify/fastai/callbacks/fp16.py

@@ -0,0 +1,122 @@
+"Callback support for half precision (fp16) training. Increases training speed."
+from ..torch_core import *
+from ..callback import *
+from ..basic_train import *
+from torch._utils import _unflatten_dense_tensors
+from torch.nn.utils import parameters_to_vector
+
+__all__ = ['MixedPrecision']
+
+def get_master(layer_groups:ModuleList, flat_master:bool=False) -> Tuple[List[List[Tensor]], List[List[Tensor]]]:
+    "Return two lists, one for the model parameters in FP16 and one for the master parameters in FP32."
+    split_params = split_no_wd_params(layer_groups)
+    model_params = [[param for param in pg if param.requires_grad] for pg in split_params]
+    if flat_master:
+        master_params = []
+        for lg in model_params:
+            if len(lg) !=0 :
+                mp = parameters_to_vector([param.data.float() for param in lg])
+                mp = torch.nn.Parameter(mp, requires_grad=True)
+                if mp.grad is None: mp.grad = mp.new(*mp.size())
+                master_params.append([mp])
+            else: master_params.append([])
+        return model_params, master_params
+    else:
+        master_params = [[param.clone().float().detach() for param in lg] for lg in model_params]
+        for mp in master_params:
+            for param in mp: param.requires_grad = True
+        return model_params, master_params
+
+def model_g2master_g(model_params:Sequence[Tensor], master_params:Sequence[Tensor], flat_master:bool=False)->None:
+    "Copy the `model_params` gradients to `master_params` for the optimizer step."
+    if flat_master:
+        for model_group,master_group in zip(model_params,master_params):
+            if len(master_group) != 0:
+                if master_group[0].grad is None: master_group[0].grad = master_group[0].data.new(*master_group[0].data.size())
+                master_group[0].grad.data.copy_(parameters_to_vector([p.grad.data.float() for p in model_group]))
+    else:
+        for model_group,master_group in zip(model_params,master_params):
+            for model, master in zip(model_group, master_group):
+                if model.grad is not None:
+                    if master.grad is None: master.grad = master.data.new(*master.data.size())
+                    master.grad.data.copy_(model.grad.data)
+                else: master.grad = None
+
+def master2model(model_params:Sequence[Tensor], master_params:Sequence[Tensor], flat_master:bool=False)->None:
+    "Copy `master_params` to `model_params`."
+    if flat_master:
+        for model_group,master_group in zip(model_params,master_params):
+            if len(model_group) != 0:
+                for model, master in zip(model_group, _unflatten_dense_tensors(master_group[0].data, model_group)):
+                    model.data.copy_(master)
+    else:
+        for model_group,master_group in zip(model_params,master_params):
+            for model, master in zip(model_group, master_group): model.data.copy_(master.data)
+
+def grad_overflow(param_group):
+    for group in param_group:
+        for p in group:
+            if p.grad is not None:
+                s = float(p.grad.data.float().sum())
+                if s == float('inf') or s == float('-inf') or s != s: return True
+    return False
+
+class MixedPrecision(LearnerCallback):
+    _order = 999 #Need to run after things that could call on_backward_begin and change the loss
+    "Callback that handles mixed-precision training."
+    def __init__(self, learn:Learner, loss_scale:float=None, max_noskip:int=1000, dynamic:bool=True, clip:float=None,
+                 flat_master:bool=False, max_scale:float=2**24):
+        super().__init__(learn)
+        self.flat_master,self.dynamic,self.max_noskip,self.clip,self.max_scale = flat_master,dynamic,max_noskip,clip,max_scale
+        self.loss_scale = ifnone(loss_scale, 2**16 if dynamic else 512)
+        self.not_min += ['model_params', 'master_params']
+        assert torch.backends.cudnn.enabled, "Mixed precision training requires cudnn."
+        self.opt = None
+
+    def on_train_begin(self, **kwargs:Any)->None:
+        "Prepare the master model."
+        #Get a copy of the model params in FP32
+        self.model_params, self.master_params = get_master(self.learn.layer_groups, self.flat_master)
+        #Changes the optimizer so that the optimization step is done in FP32.
+        new_opt = self.learn.opt.new_with_params(self.master_params)
+        if self.opt is not None:
+            self.opt.lr,self.opt.wd = self.learn.opt.lr,self.learn.opt.wd
+            new_opt.load_state_dict(self.opt)
+        self.learn.opt.opt = new_opt.opt
+        self.noskip = 0
+
+    def on_loss_begin(self, last_output:Tensor, **kwargs:Any) -> Tensor:
+        "Convert half precision output to FP32 to avoid reduction overflow."
+        return {'last_output': to_float(last_output)}
+
+    def on_backward_begin(self, last_loss:Rank0Tensor, **kwargs:Any) -> Rank0Tensor:
+        "Scale gradients up by `self.loss_scale` to prevent underflow."
+        #To avoid gradient underflow, we scale the gradients
+        ret_loss = last_loss * self.loss_scale
+        return {'last_loss': ret_loss}
+
+    def on_backward_end(self, **kwargs:Any)->None:
+        "Convert the gradients back to FP32 and divide them by the scale."
+        if self.dynamic and grad_overflow(self.model_params) and self.loss_scale > 1:
+            self.loss_scale /= 2
+            self.noskip = 0
+            #The step will be skipped since we don't update the master grads so they are all None or zero
+        else:
+            model_g2master_g(self.model_params, self.master_params, self.flat_master)
+            for group in self.master_params:
+                for param in group:
+                    if param.grad is not None: param.grad.div_(self.loss_scale)
+            if self.clip is not None:
+                for group in self.master_params: nn.utils.clip_grad_norm_(group, self.clip)
+            if not self.dynamic: return
+            self.noskip += 1
+            if self.noskip >= self.max_noskip and self.loss_scale < self.max_scale:
+                self.loss_scale *= 2
+                self.noskip = 0
+
+    def on_step_end(self, **kwargs:Any)->None:
+        "Update the params from master to model and zero grad."
+        #Zeros the gradients of the model since the optimizer is disconnected.
+        self.learn.model.zero_grad()
+        #Update the params from master to model.
+        master2model(self.model_params, self.master_params, self.flat_master)

DeOldify/fastai/callbacks/general_sched.py

@@ -0,0 +1,46 @@
+from ..core import *
+from ..callback import *
+from ..basic_train import Learner, LearnerCallback
+
+__all__ = ['GeneralScheduler', 'TrainingPhase']
+
+@dataclass
+class TrainingPhase():
+    "Schedule hyper-parameters for a phase of `length` iterations."
+    length:int
+    
+    def __post_init__(self): self.scheds = dict()
+    def schedule_hp(self, name, vals, anneal=None):
+        "Adds a schedule for `name` between `vals` using `anneal`."
+        self.scheds[name] = Scheduler(vals, self.length, anneal)
+        return self
+
+class GeneralScheduler(LearnerCallback):
+    "Schedule multiple `TrainingPhase` for a `Learner`."
+    def __init__(self, learn:Learner, phases:Collection[TrainingPhase], start_epoch:int=None):
+        super().__init__(learn)
+        self.phases,self.start_epoch = phases,start_epoch
+
+    def on_train_begin(self, epoch:int, **kwargs:Any)->None:
+        "Initialize the schedulers for training."
+        res = {'epoch':self.start_epoch} if self.start_epoch is not None else None
+        self.start_epoch = ifnone(self.start_epoch, epoch)
+        self.scheds = [p.scheds for p in self.phases]
+        self.opt = self.learn.opt
+        for k,v in self.scheds[0].items(): 
+            v.restart()
+            self.opt.set_stat(k, v.start)
+        self.idx_s = 0
+        return res
+    
+    def jump_to_epoch(self, epoch:int)->None:
+        for _ in range(len(self.learn.data.train_dl) * epoch):
+            self.on_batch_end(True)
+
+    def on_batch_end(self, train, **kwargs:Any)->None:
+        "Take a step in lr,mom sched, start next stepper when the current one is complete."
+        if train:
+            if self.idx_s >= len(self.scheds): return {'stop_training': True, 'stop_epoch': True}
+            sched = self.scheds[self.idx_s]
+            for k,v in sched.items(): self.opt.set_stat(k, v.step())
+            if list(sched.values())[0].is_done: self.idx_s += 1

DeOldify/fastai/callbacks/hooks.py

@@ -0,0 +1,200 @@
+"Hooks provide extensibility at the model level."
+from ..torch_core import *
+from ..callback import *
+from ..basic_train import *
+from ..basic_data import *
+
+__all__ = ['ActivationStats', 'Hook', 'HookCallback', 'Hooks', 'hook_output', 'hook_outputs',
+           'model_sizes', 'num_features_model', 'model_summary', 'dummy_eval', 'dummy_batch']
+
+class Hook():
+    "Create a hook on `m` with `hook_func`."
+    def __init__(self, m:nn.Module, hook_func:HookFunc, is_forward:bool=True, detach:bool=True):
+        self.hook_func,self.detach,self.stored = hook_func,detach,None
+        f = m.register_forward_hook if is_forward else m.register_backward_hook
+        self.hook = f(self.hook_fn)
+        self.removed = False
+
+    def hook_fn(self, module:nn.Module, input:Tensors, output:Tensors):
+        "Applies `hook_func` to `module`, `input`, `output`."
+        if self.detach:
+            input  = (o.detach() for o in input ) if is_listy(input ) else input.detach()
+            output = (o.detach() for o in output) if is_listy(output) else output.detach()
+        self.stored = self.hook_func(module, input, output)
+
+    def remove(self):
+        "Remove the hook from the model."
+        if not self.removed:
+            self.hook.remove()
+            self.removed=True
+
+    def __enter__(self, *args): return self
+    def __exit__(self, *args): self.remove()
+
+class Hooks():
+    "Create several hooks on the modules in `ms` with `hook_func`."
+    def __init__(self, ms:Collection[nn.Module], hook_func:HookFunc, is_forward:bool=True, detach:bool=True):
+        self.hooks = [Hook(m, hook_func, is_forward, detach) for m in ms]
+
+    def __getitem__(self,i:int)->Hook: return self.hooks[i]
+    def __len__(self)->int: return len(self.hooks)
+    def __iter__(self): return iter(self.hooks)
+    @property
+    def stored(self): return [o.stored for o in self]
+
+    def remove(self):
+        "Remove the hooks from the model."
+        for h in self.hooks: h.remove()
+
+    def __enter__(self, *args): return self
+    def __exit__ (self, *args): self.remove()
+
+def _hook_inner(m,i,o): return o if isinstance(o,Tensor) else o if is_listy(o) else list(o)
+
+def hook_output (module:nn.Module, detach:bool=True, grad:bool=False)->Hook:
+    "Return a `Hook` that stores activations of `module` in `self.stored`"
+    return Hook(module, _hook_inner, detach=detach, is_forward=not grad)
+
+def hook_outputs(modules:Collection[nn.Module], detach:bool=True, grad:bool=False)->Hooks:
+    "Return `Hooks` that store activations of all `modules` in `self.stored`"
+    return Hooks(modules, _hook_inner, detach=detach, is_forward=not grad)
+
+class HookCallback(LearnerCallback):
+    "Callback that can be used to register hooks on `modules`. Implement the corresponding function in `self.hook`."
+    def __init__(self, learn:Learner, modules:Sequence[nn.Module]=None, do_remove:bool=True):
+        super().__init__(learn)
+        self.modules,self.do_remove = modules,do_remove
+
+    def on_train_begin(self, **kwargs):
+        "Register the `Hooks` on `self.modules`."
+        if not self.modules:
+            self.modules = [m for m in flatten_model(self.learn.model)
+                            if hasattr(m, 'weight')]
+        self.hooks = Hooks(self.modules, self.hook)
+
+    def on_train_end(self, **kwargs):
+        "Remove the `Hooks`."
+        if self.do_remove: self.remove()
+
+    def remove(self): 
+        if getattr(self, 'hooks', None): self.hooks.remove()
+    def __del__(self): self.remove()
+
+class ActivationStats(HookCallback):
+    "Callback that record the mean and std of activations."
+
+    def on_train_begin(self, **kwargs):
+        "Initialize stats."
+        super().on_train_begin(**kwargs)
+        self.stats = []
+
+    def hook(self, m:nn.Module, i:Tensors, o:Tensors)->Tuple[Rank0Tensor,Rank0Tensor]:
+        "Take the mean and std of `o`."
+        return o.mean().item(),o.std().item()
+    def on_batch_end(self, train, **kwargs):
+        "Take the stored results and puts it in `self.stats`"
+        if train: self.stats.append(self.hooks.stored)
+    def on_train_end(self, **kwargs):
+        "Polish the final result."
+        super().on_train_end(**kwargs)
+        self.stats = tensor(self.stats).permute(2,1,0)
+
+def dummy_batch(m: nn.Module, size:tuple=(64,64))->Tensor:
+    "Create a dummy batch to go through `m` with `size`."
+    ch_in = in_channels(m)
+    return one_param(m).new(1, ch_in, *size).requires_grad_(False).uniform_(-1.,1.)
+
+def dummy_eval(m:nn.Module, size:tuple=(64,64)):
+    "Pass a `dummy_batch` in evaluation mode in `m` with `size`."
+    m.eval()
+    return m(dummy_batch(m, size))
+    #return m.eval()(dummy_batch(m, size))
+
+def model_sizes(m:nn.Module, size:tuple=(64,64))->Tuple[Sizes,Tensor,Hooks]:
+    "Pass a dummy input through the model `m` to get the various sizes of activations."
+    with hook_outputs(m) as hooks:
+        x = dummy_eval(m, size)
+        return [o.stored.shape for o in hooks]
+
+def num_features_model(m:nn.Module)->int:
+    "Return the number of output features for `model`."
+    sz = 64
+    while True:
+        try: return model_sizes(m, size=(sz,sz))[-1][1]
+        except Exception as e:
+            sz *= 2
+            if sz > 2048: raise
+
+def total_params(m:nn.Module)->int:
+    params, trainable = 0, False
+    if hasattr(m, "weight") and hasattr(m.weight, "size"):
+         params += m.weight.numel()
+         trainable = m.weight.requires_grad
+    if hasattr(m, "bias") and hasattr(m.bias, "size"): params += m.bias.numel()
+    return params, trainable
+
+def hook_params(modules:Collection[nn.Module])->Hooks:
+    return Hooks(modules, lambda m, i, o: total_params(m))
+
+def params_size(m: Union[nn.Module,Learner], size: tuple = (3, 64, 64))->Tuple[Sizes, Tensor, Hooks]:
+    "Pass a dummy input through the model to get the various sizes. Returns (res,x,hooks) if `full`"
+    if isinstance(m, Learner):
+        if m.data.is_empty:
+            raise Exception("This is an empty `Learner` and `Learner.summary` requires some data to pass through the model.")
+        ds_type = DatasetType.Train if m.data.train_dl else (DatasetType.Valid if m.data.valid_dl else DatasetType.Test)
+        x = m.data.one_batch(ds_type=ds_type, detach=False, denorm=False)[0]
+        x = [o[:1] for o in x]  if is_listy(x) else x[:1]
+        m = m.model
+    elif isinstance(m, nn.Module): x = next(m.parameters()).new(1, *size)
+    else: raise TypeError('You should either pass in a Learner or nn.Module')
+    with hook_outputs(flatten_model(m)) as hook_o:
+        with hook_params(flatten_model(m))as hook_p:
+            x = m.eval()(*x) if is_listy(x) else m.eval()(x)
+            output_size = [((o.stored.shape[1:]) if o.stored is not None else None) for o in hook_o]
+            params = [(o.stored if o.stored is not None else (None,None)) for o in hook_p]
+    params, trainables = map(list,zip(*params))
+    return output_size, params, trainables
+
+def get_layer_name(layer:nn.Module)->str:
+    return str(layer.__class__).split(".")[-1].split("'")[0]
+
+def layers_info(m:Collection[nn.Module]) -> Collection[namedtuple]:
+    func = lambda m:list(map(get_layer_name, flatten_model(m)))
+    layers_names = func(m.model) if isinstance(m, Learner) else func(m)
+    layers_sizes, layers_params, layers_trainable = params_size(m)
+    layer_info = namedtuple('Layer_Information', ['Layer', 'OutputSize', 'Params', 'Trainable'])
+    return list(map(layer_info, layers_names, layers_sizes, layers_params, layers_trainable))
+
+def model_summary(m:Learner, n:int=70):
+    "Print a summary of `m` using a output text width of `n` chars"
+    info = layers_info(m)
+    header = ["Layer (type)", "Output Shape", "Param #", "Trainable"]
+    res = m.model.__class__.__name__ + "\n"
+    res += "=" * n + "\n"
+    res += f"{header[0]:<20} {header[1]:<20} {header[2]:<10} {header[3]:<10}\n"
+    res += "=" * n + "\n"
+    total_params = 0
+    total_trainable_params = 0
+    for layer, size, params, trainable in info:
+        if size is None: continue
+        total_params += int(params)
+        total_trainable_params += int(params) * trainable
+        size, trainable = str(list(size)), str(trainable)
+        res += f"{layer:<20} {size:<20} {int(params):<10,} {trainable:<10}\n"
+        res += "_" * n + "\n"
+    res += f"\nTotal params: {total_params:,}\n"
+    res += f"Total trainable params: {total_trainable_params:,}\n"
+    res += f"Total non-trainable params: {total_params - total_trainable_params:,}\n"
+           
+    res += f"Optimized with {str(m.opt_func)[25:-1].replace('>', '')}\n"
+    if m.true_wd: res += f"Using true weight decay as discussed in https://www.fast.ai/2018/07/02/adam-weight-decay/ \n"
+    if "wd" in str(m.opt_func) or "weight_decay" in str(m.opt_func): res += f"\x1b[1;31m Specifying weight decay in the optimizer has no effect, Learner will overwrite \x1b[0m \n"
+    if "lr" in str(m.opt_func) or "learning_rate" in str(m.opt_func): res += f"\x1b[1;31m Specifying lr in the optimizer has no effect, pass it to fit or the defaults.lr will apply \x1b[0m \n" 
+    res += f"Loss function : {m.loss_func.__class__.__name__}\n"
+    res += "=" * n + "\n"
+    res += "Callbacks functions applied \n"
+    res += "\n".join([f"    {cbs.__class__.__name__}" for cbs in m.callbacks])
+
+    return PrettyString(res)
+
+Learner.summary = model_summary

DeOldify/fastai/callbacks/loss_metrics.py

@@ -0,0 +1,34 @@
+from ..torch_core import *
+from ..callback import *
+from ..basic_train import Learner, LearnerCallback
+
+__all__ = ['LossMetrics']
+
+class LossMetrics(LearnerCallback):
+    "Add `loss_func.metrics` to metrics named by `loss_func.metric_names`"
+    _order = -20 #Needs to run before the recorder
+
+    def on_train_begin(self, **kwargs):
+        "Add the metrics names to the `Recorder`."
+        self.names = ifnone(self.learn.loss_func.metric_names, [])
+        if not self.names: warn('LossMetrics requested but no loss_func.metric_names provided')
+        self.learn.recorder.add_metric_names(self.names)
+
+    def on_epoch_begin(self, **kwargs):
+        "Initialize the metrics for this epoch."
+        self.metrics = {name:0. for name in self.names}
+        self.nums = 0
+
+    def on_batch_end(self, last_target, train, **kwargs):
+        "Update the metrics if not `train`"
+        if train: return
+        bs = last_target.size(0)
+        for name in self.names:
+            self.metrics[name] += bs * self.learn.loss_func.metrics[name].detach().cpu()
+        self.nums += bs
+
+    def on_epoch_end(self, last_metrics, **kwargs):
+        "Finish the computation and sends the result to the Recorder."
+        if not self.nums: return
+        metrics = [self.metrics[name]/self.nums for name in self.names]
+        return {'last_metrics': last_metrics+metrics}

DeOldify/fastai/callbacks/lr_finder.py

@@ -0,0 +1,39 @@
+"Tools to help find the optimal learning rate for training"
+from ..torch_core import *
+from ..basic_data import DataBunch
+from ..callback import *
+from ..basic_train import Learner, LearnerCallback
+
+__all__ = ['LRFinder']
+
+class LRFinder(LearnerCallback):
+    "Causes `learn` to go on a mock training from `start_lr` to `end_lr` for `num_it` iterations."
+    def __init__(self, learn:Learner, start_lr:float=1e-7, end_lr:float=10, num_it:int=100, stop_div:bool=True):
+        super().__init__(learn)
+        self.data,self.stop_div = learn.data,stop_div
+        self.sched = Scheduler((start_lr, end_lr), num_it, annealing_exp)
+
+    def on_train_begin(self, pbar, **kwargs:Any)->None:
+        "Initialize optimizer and learner hyperparameters."
+        setattr(pbar, 'clean_on_interrupt', True)
+        self.learn.save('tmp')
+        self.opt = self.learn.opt
+        self.opt.lr = self.sched.start
+        self.stop,self.best_loss = False,0.
+        return {'skip_validate': True}
+
+    def on_batch_end(self, iteration:int, smooth_loss:TensorOrNumber, **kwargs:Any)->None:
+        "Determine if loss has runaway and we should stop."
+        if iteration==0 or smooth_loss < self.best_loss: self.best_loss = smooth_loss
+        self.opt.lr = self.sched.step()
+        if self.sched.is_done or (self.stop_div and (smooth_loss > 4*self.best_loss or torch.isnan(smooth_loss))):
+            #We use the smoothed loss to decide on the stopping since it's less shaky.
+            return {'stop_epoch': True, 'stop_training': True}
+
+    def on_train_end(self, **kwargs:Any)->None:
+        "Cleanup learn model weights disturbed during LRFinder exploration."
+        self.learn.load('tmp', purge=False)
+        if hasattr(self.learn.model, 'reset'): self.learn.model.reset()
+        for cb in self.callbacks:
+            if hasattr(cb, 'reset'): cb.reset()
+        print('LR Finder is complete, type {learner_name}.recorder.plot() to see the graph.')

DeOldify/fastai/callbacks/mem.py

@@ -0,0 +1,65 @@
+" Memory profiling callbacks "
+
+import tracemalloc, threading, torch, time
+from ..utils.mem import *
+from ..basic_train import *
+from ..torch_core import *
+from ..utils.pynvml_gate import *
+
+if use_gpu: pynvml = load_pynvml_env()
+
+class PeakMemMetric(LearnerCallback):
+    "Callback that measures used and peaked general and GPU memory."
+
+    _order=-20 # Needs to run before the recorder
+
+    def __init__(self, learn:Learner):
+        super().__init__(learn)
+        assert torch.cuda.is_available(), "pytorch CUDA is required"
+        preload_pytorch()
+
+    def peak_monitor_start(self):
+        self.peak_monitoring = True
+
+        # start RAM tracing
+        tracemalloc.start()
+
+        # this thread samples RAM usage as long as the current epoch of the fit loop is running
+        peak_monitor_thread = threading.Thread(target=self.peak_monitor_func)
+        peak_monitor_thread.daemon = True
+        peak_monitor_thread.start()
+
+    def peak_monitor_stop(self):
+        tracemalloc.stop()
+        self.peak_monitoring = False
+
+    def peak_monitor_func(self):
+        self.gpu_mem_used_peak = -1
+
+        gpu_id = torch.cuda.current_device()
+        gpu_handle = pynvml.nvmlDeviceGetHandleByIndex(gpu_id)
+
+        while True:
+            gpu_mem_used = gpu_mem_get_used_fast(gpu_handle)
+            self.gpu_mem_used_peak = max(gpu_mem_used, self.gpu_mem_used_peak)
+            if not self.peak_monitoring: break
+            time.sleep(0.001) # 1msec
+
+    def on_train_begin(self, **kwargs):
+        self.learn.recorder.add_metric_names(['cpu used',  'peak', 'gpu used',  'peak'])
+
+    def on_epoch_begin(self, **kwargs):
+        self.peak_monitor_start()
+        self.gpu_before = gpu_mem_get_used_no_cache()
+
+    def on_epoch_end(self, last_metrics, **kwargs):
+        cpu_used, cpu_peak =  list(map(lambda x: int(x/2**20), tracemalloc.get_traced_memory()))
+        self.peak_monitor_stop()
+        gpu_used = gpu_mem_get_used_no_cache() - self.gpu_before
+        gpu_peak = self.gpu_mem_used_peak      - self.gpu_before
+        # can be negative, due to unreliable peak monitor thread
+        if gpu_peak < 0:   gpu_peak = 0
+        # since we want the overhead only, subtract delta used if it's positive
+        elif gpu_used > 0: gpu_peak -= gpu_used
+        # The numbers are deltas in MBs (beginning of the epoch and the end)
+        return add_metrics(last_metrics, [cpu_used, cpu_peak, gpu_used, gpu_peak])

DeOldify/fastai/callbacks/misc.py

@@ -0,0 +1,12 @@
+" Miscellaneous callbacks "
+
+from fastai.callback import Callback
+
+class StopAfterNBatches(Callback):
+    "Stop training after n batches of the first epoch."
+    def __init__(self, n_batches:int=2):
+        self.stop,self.n_batches = False,n_batches-1 # iteration starts from 0
+
+    def on_batch_end(self, iteration, **kwargs):
+        if iteration == self.n_batches:
+            return {'stop_epoch': True, 'stop_training': True, 'skip_validate': True}

DeOldify/fastai/callbacks/mixup.py

@@ -0,0 +1,67 @@
+"Implements [mixup](https://arxiv.org/abs/1710.09412) training method"
+from ..torch_core import *
+from ..callback import *
+from ..basic_train import Learner, LearnerCallback
+
+class MixUpCallback(LearnerCallback):
+    "Callback that creates the mixed-up input and target."
+    def __init__(self, learn:Learner, alpha:float=0.4, stack_x:bool=False, stack_y:bool=True):
+        super().__init__(learn)
+        self.alpha,self.stack_x,self.stack_y = alpha,stack_x,stack_y
+    
+    def on_train_begin(self, **kwargs):
+        if self.stack_y: self.learn.loss_func = MixUpLoss(self.learn.loss_func)
+        
+    def on_batch_begin(self, last_input, last_target, train, **kwargs):
+        "Applies mixup to `last_input` and `last_target` if `train`."
+        if not train: return
+        lambd = np.random.beta(self.alpha, self.alpha, last_target.size(0))
+        lambd = np.concatenate([lambd[:,None], 1-lambd[:,None]], 1).max(1)
+        lambd = last_input.new(lambd)
+        shuffle = torch.randperm(last_target.size(0)).to(last_input.device)
+        x1, y1 = last_input[shuffle], last_target[shuffle]
+        if self.stack_x:
+            new_input = [last_input, last_input[shuffle], lambd]
+        else: 
+            out_shape = [lambd.size(0)] + [1 for _ in range(len(x1.shape) - 1)]
+            new_input = (last_input * lambd.view(out_shape) + x1 * (1-lambd).view(out_shape))
+        if self.stack_y:
+            new_target = torch.cat([last_target[:,None].float(), y1[:,None].float(), lambd[:,None].float()], 1)
+        else:
+            if len(last_target.shape) == 2:
+                lambd = lambd.unsqueeze(1).float()
+            new_target = last_target.float() * lambd + y1.float() * (1-lambd)
+        return {'last_input': new_input, 'last_target': new_target}  
+    
+    def on_train_end(self, **kwargs):
+        if self.stack_y: self.learn.loss_func = self.learn.loss_func.get_old()
+        
+
+class MixUpLoss(Module):
+    "Adapt the loss function `crit` to go with mixup."
+    
+    def __init__(self, crit, reduction='mean'):
+        super().__init__()
+        if hasattr(crit, 'reduction'): 
+            self.crit = crit
+            self.old_red = crit.reduction
+            setattr(self.crit, 'reduction', 'none')
+        else: 
+            self.crit = partial(crit, reduction='none')
+            self.old_crit = crit
+        self.reduction = reduction
+        
+    def forward(self, output, target):
+        if len(target.size()) == 2:
+            loss1, loss2 = self.crit(output,target[:,0].long()), self.crit(output,target[:,1].long())
+            d = (loss1 * target[:,2] + loss2 * (1-target[:,2])).mean()
+        else:  d = self.crit(output, target)
+        if self.reduction == 'mean': return d.mean()
+        elif self.reduction == 'sum':            return d.sum()
+        return d
+    
+    def get_old(self):
+        if hasattr(self, 'old_crit'):  return self.old_crit
+        elif hasattr(self, 'old_red'): 
+            setattr(self.crit, 'reduction', self.old_red)
+            return self.crit

DeOldify/fastai/callbacks/mlflow.py

@@ -0,0 +1,36 @@
+"A `Callback` that saves tracked metrics and notebook file into MLflow server."
+from ..torch_core import *
+from ..callback import *
+from ..basic_train import Learner, LearnerCallback
+#This is an optional dependency in fastai.  Must install separately.
+try: import mlflow
+except: print("To use this tracker, please run 'pip install mlflow'")
+
+class MLFlowTracker(LearnerCallback):
+    "A `TrackerCallback` that tracks the loss and metrics into MLFlow"
+    def __init__(self, learn:Learner, exp_name: str, params: dict, nb_path: str, uri: str = "http://localhost:5000"):
+        super().__init__(learn)
+        self.learn,self.exp_name,self.params,self.nb_path,self.uri = learn,exp_name,params,nb_path,uri
+        self.metrics_names = ['train_loss', 'valid_loss'] + [o.__name__ for o in learn.metrics]
+
+    def on_train_begin(self, **kwargs: Any) -> None:
+        "Prepare MLflow experiment and log params"
+        self.client = mlflow.tracking.MlflowClient(self.uri)
+        exp = self.client.get_experiment_by_name(self.exp_name)
+        self.exp_id = self.client.create_experiment(self.exp_name) if exp is None else exp.experiment_id
+        run = self.client.create_run(experiment_id=self.exp_id)
+        self.run = run.info.run_uuid
+        for k,v in self.params.items():
+            self.client.log_param(run_id=self.run, key=k, value=v)
+
+    def on_epoch_end(self, epoch, **kwargs:Any)->None:
+        "Send loss and metrics values to MLFlow after each epoch"
+        if kwargs['smooth_loss'] is None or kwargs["last_metrics"] is None: return
+        metrics = [kwargs['smooth_loss']] + kwargs["last_metrics"]
+        for name, val in zip(self.metrics_names, metrics):
+            self.client.log_metric(self.run, name, np.float(val), step=epoch)
+        
+    def on_train_end(self, **kwargs: Any) -> None:  
+        "Store the notebook and stop run"
+        self.client.log_artifact(run_id=self.run, local_path=self.nb_path)
+        self.client.set_terminated(run_id=self.run)

DeOldify/fastai/callbacks/one_cycle.py

@@ -0,0 +1,58 @@
+"Supports 1-Cycle style training"
+from ..core import *
+from ..callback import *
+from ..basic_train import Learner,LearnerCallback
+
+__all__ = ['OneCycleScheduler']
+
+class OneCycleScheduler(LearnerCallback):
+    "Manage 1-Cycle style training as outlined in Leslie Smith's [paper](https://arxiv.org/pdf/1803.09820.pdf)."
+    def __init__(self, learn:Learner, lr_max:float, moms:Floats=(0.95,0.85), div_factor:float=25., pct_start:float=0.3,
+                 final_div:float=None, tot_epochs:int=None, start_epoch:int=None):
+        super().__init__(learn)
+        self.lr_max,self.div_factor,self.pct_start,self.final_div = lr_max,div_factor,pct_start,final_div
+        if self.final_div is None: self.final_div = div_factor*1e4
+        self.moms=tuple(listify(moms,2))
+        if is_listy(self.lr_max): self.lr_max = np.array(self.lr_max)
+        self.start_epoch, self.tot_epochs = start_epoch, tot_epochs
+
+    def steps(self, *steps_cfg:StartOptEnd):
+        "Build anneal schedule for all of the parameters."
+        return [Scheduler(step, n_iter, func=func)
+                for (step,(n_iter,func)) in zip(steps_cfg, self.phases)]
+
+    def on_train_begin(self, n_epochs:int, epoch:int, **kwargs:Any)->None:
+        "Initialize our optimization params based on our annealing schedule."
+        res = {'epoch':self.start_epoch} if self.start_epoch is not None else None
+        self.start_epoch = ifnone(self.start_epoch, epoch)
+        self.tot_epochs = ifnone(self.tot_epochs, n_epochs)
+        n = len(self.learn.data.train_dl) * self.tot_epochs
+        a1 = int(n * self.pct_start)
+        a2 = n-a1
+        self.phases = ((a1, annealing_cos), (a2, annealing_cos))
+        low_lr = self.lr_max/self.div_factor
+        self.lr_scheds = self.steps((low_lr, self.lr_max), (self.lr_max, self.lr_max/self.final_div))
+        self.mom_scheds = self.steps(self.moms, (self.moms[1], self.moms[0]))
+        self.opt = self.learn.opt
+        self.opt.lr,self.opt.mom = self.lr_scheds[0].start,self.mom_scheds[0].start
+        self.idx_s = 0
+        return res
+    
+    def jump_to_epoch(self, epoch:int)->None:
+        for _ in range(len(self.learn.data.train_dl) * epoch):
+            self.on_batch_end(True)
+
+    def on_batch_end(self, train, **kwargs:Any)->None:
+        "Take one step forward on the annealing schedule for the optim params."
+        if train:
+            if self.idx_s >= len(self.lr_scheds): return {'stop_training': True, 'stop_epoch': True}
+            self.opt.lr = self.lr_scheds[self.idx_s].step()
+            self.opt.mom = self.mom_scheds[self.idx_s].step()
+            # when the current schedule is complete we move onto the next
+            # schedule. (in 1-cycle there are two schedules)
+            if self.lr_scheds[self.idx_s].is_done:
+                self.idx_s += 1
+
+    def on_epoch_end(self, epoch, **kwargs:Any)->None:
+        "Tell Learner to stop if the cycle is finished."
+        if epoch > self.tot_epochs: return {'stop_training': True}

DeOldify/fastai/callbacks/oversampling.py

@@ -0,0 +1,22 @@
+from ..torch_core import *
+from ..basic_data import DataBunch
+from ..callback import *
+from ..basic_train import Learner,LearnerCallback
+from torch.utils.data.sampler import WeightedRandomSampler
+
+__all__ = ['OverSamplingCallback']
+
+
+
+class OverSamplingCallback(LearnerCallback):
+    def __init__(self,learn:Learner,weights:torch.Tensor=None):
+        super().__init__(learn)
+        self.labels = self.learn.data.train_dl.dataset.y.items
+        _, counts = np.unique(self.labels,return_counts=True)
+        self.weights = (weights if weights is not None else
+                        torch.DoubleTensor((1/counts)[self.labels]))
+        self.label_counts = np.bincount([self.learn.data.train_dl.dataset.y[i].data for i in range(len(self.learn.data.train_dl.dataset))])
+        self.total_len_oversample = int(self.learn.data.c*np.max(self.label_counts))
+        
+    def on_train_begin(self, **kwargs):
+        self.learn.data.train_dl.dl.batch_sampler = BatchSampler(WeightedRandomSampler(self.weights,self.total_len_oversample), self.learn.data.train_dl.batch_size,False)

DeOldify/fastai/callbacks/rnn.py

@@ -0,0 +1,31 @@
+"Regroups lr adjustment to seq_len, AR and TAR"
+from ..torch_core import *
+from ..callback import *
+from ..basic_train import Learner, LearnerCallback
+
+__all__ = ['RNNTrainer']
+
+class RNNTrainer(LearnerCallback):
+    "`Callback` that regroups lr adjustment to seq_len, AR and TAR."
+    def __init__(self, learn:Learner, alpha:float=0., beta:float=0.):
+        super().__init__(learn)
+        self.not_min += ['raw_out', 'out']
+        self.alpha,self.beta = alpha,beta
+        
+    def on_epoch_begin(self, **kwargs):
+        "Reset the hidden state of the model."
+        self.learn.model.reset()
+
+    def on_loss_begin(self, last_output:Tuple[Tensor,Tensor,Tensor], **kwargs):
+        "Save the extra outputs for later and only returns the true output."
+        self.raw_out,self.out = last_output[1],last_output[2]
+        return {'last_output': last_output[0]}
+
+    def on_backward_begin(self, last_loss:Rank0Tensor, last_input:Tensor, **kwargs):
+        "Apply AR and TAR to `last_loss`."
+        #AR and TAR
+        if self.alpha != 0.:  last_loss += self.alpha * self.out[-1].float().pow(2).mean()
+        if self.beta != 0.:
+            h = self.raw_out[-1]
+            if len(h)>1: last_loss += self.beta * (h[:,1:] - h[:,:-1]).float().pow(2).mean()
+        return {'last_loss': last_loss}

DeOldify/fastai/callbacks/tensorboard.py

@@ -0,0 +1,427 @@
+"Provides convenient callbacks for Learners that write model images, metrics/losses, stats and histograms to Tensorboard"
+from ..basic_train import Learner
+from ..basic_data import DatasetType, DataBunch
+from ..vision import Image
+from ..vision.gan import GANLearner
+from ..callbacks import LearnerCallback
+from ..core import *
+from ..torch_core import *
+from threading import Thread, Event
+from time import sleep
+from queue import Queue
+import statistics
+import torchvision.utils as vutils
+from abc import ABC
+#This is an optional dependency in fastai.  Must install separately.
+try: from tensorboardX import SummaryWriter
+except: print("To use this tracker, please run 'pip install tensorboardx'. Also you must have Tensorboard running to see results")
+
+__all__=['LearnerTensorboardWriter', 'GANTensorboardWriter', 'ImageGenTensorboardWriter']
+
+#---Example usage (applies to any of the callbacks)--- 
+# proj_id = 'Colorize'
+# tboard_path = Path('data/tensorboard/' + proj_id)
+# learn.callback_fns.append(partial(GANTensorboardWriter, base_dir=tboard_path, name='GanLearner'))
+
+class LearnerTensorboardWriter(LearnerCallback):
+    "Broadly useful callback for Learners that writes to Tensorboard.  Writes model histograms, losses/metrics, and gradient stats."
+    def __init__(self, learn:Learner, base_dir:Path, name:str, loss_iters:int=25, hist_iters:int=500, stats_iters:int=100):
+        super().__init__(learn=learn)
+        self.base_dir,self.name,self.loss_iters,self.hist_iters,self.stats_iters  = base_dir,name,loss_iters,hist_iters,stats_iters
+        log_dir = base_dir/name
+        self.tbwriter = SummaryWriter(str(log_dir))
+        self.hist_writer = HistogramTBWriter()
+        self.stats_writer = ModelStatsTBWriter()
+        #self.graph_writer = GraphTBWriter()
+        self.data = None
+        self.metrics_root = '/metrics/'
+        self._update_batches_if_needed()
+
+    def _get_new_batch(self, ds_type:DatasetType)->Collection[Tensor]:
+        "Retrieves new batch of DatasetType, and detaches it."
+        return self.learn.data.one_batch(ds_type=ds_type, detach=True, denorm=False, cpu=False)
+
+    def _update_batches_if_needed(self)->None:
+        "one_batch function is extremely slow with large datasets.  This is caching the result as an optimization."
+        if self.learn.data.valid_dl is None: return # Running learning rate finder, so return
+        update_batches = self.data is not self.learn.data
+        if not update_batches: return
+        self.data = self.learn.data
+        self.trn_batch = self._get_new_batch(ds_type=DatasetType.Train)
+        self.val_batch = self._get_new_batch(ds_type=DatasetType.Valid)
+
+    def _write_model_stats(self, iteration:int)->None:
+        "Writes gradient statistics to Tensorboard."
+        self.stats_writer.write(model=self.learn.model, iteration=iteration, tbwriter=self.tbwriter)
+
+    def _write_training_loss(self, iteration:int, last_loss:Tensor)->None:
+        "Writes training loss to Tensorboard."
+        scalar_value = to_np(last_loss)
+        tag = self.metrics_root + 'train_loss'
+        self.tbwriter.add_scalar(tag=tag, scalar_value=scalar_value, global_step=iteration)
+
+    def _write_weight_histograms(self, iteration:int)->None:
+        "Writes model weight histograms to Tensorboard."
+        self.hist_writer.write(model=self.learn.model, iteration=iteration, tbwriter=self.tbwriter)
+
+    def _write_scalar(self, name:str, scalar_value, iteration:int)->None:
+        "Writes single scalar value to Tensorboard."
+        tag = self.metrics_root + name
+        self.tbwriter.add_scalar(tag=tag, scalar_value=scalar_value, global_step=iteration)
+
+    #TODO:  Relying on a specific hardcoded start_idx here isn't great.  Is there a better solution?
+    def _write_metrics(self, iteration:int, last_metrics:MetricsList, start_idx:int=2)->None:
+        "Writes training metrics to Tensorboard."
+        recorder = self.learn.recorder
+        for i, name in enumerate(recorder.names[start_idx:]):
+            if last_metrics is None or len(last_metrics) < i+1: return
+            scalar_value = last_metrics[i]
+            self._write_scalar(name=name, scalar_value=scalar_value, iteration=iteration)
+
+    def on_train_begin(self, **kwargs: Any) -> None:
+        #self.graph_writer.write(model=self.learn.model, tbwriter=self.tbwriter,
+                                #input_to_model=next(iter(self.learn.data.dl(DatasetType.Single)))[0])
+        return
+
+    def on_batch_end(self, last_loss:Tensor, iteration:int, **kwargs)->None:
+        "Callback function that writes batch end appropriate data to Tensorboard."
+        if iteration == 0: return
+        self._update_batches_if_needed()
+        if iteration % self.loss_iters == 0: self._write_training_loss(iteration=iteration, last_loss=last_loss)
+        if iteration % self.hist_iters == 0: self._write_weight_histograms(iteration=iteration)
+
+    # Doing stuff here that requires gradient info, because they get zeroed out afterwards in training loop
+    def on_backward_end(self, iteration:int, **kwargs)->None:
+        "Callback function that writes backward end appropriate data to Tensorboard."
+        if iteration == 0: return
+        self._update_batches_if_needed()
+        if iteration % self.stats_iters == 0: self._write_model_stats(iteration=iteration)
+
+    def on_epoch_end(self, last_metrics:MetricsList, iteration:int, **kwargs)->None:
+        "Callback function that writes epoch end appropriate data to Tensorboard."
+        self._write_metrics(iteration=iteration, last_metrics=last_metrics)
+
+# TODO:  We're overriding almost everything here.  Seems like a good idea to question that ("is a" vs "has a")
+class GANTensorboardWriter(LearnerTensorboardWriter):
+    "Callback for GANLearners that writes to Tensorboard.  Extends LearnerTensorboardWriter and adds output image writes."
+    def __init__(self, learn:GANLearner, base_dir:Path, name:str, loss_iters:int=25, hist_iters:int=500, 
+                stats_iters:int=100, visual_iters:int=100):
+        super().__init__(learn=learn, base_dir=base_dir, name=name, loss_iters=loss_iters, hist_iters=hist_iters, stats_iters=stats_iters)
+        self.visual_iters = visual_iters
+        self.img_gen_vis = ImageTBWriter()
+        self.gen_stats_updated = True
+        self.crit_stats_updated = True
+
+    def _write_weight_histograms(self, iteration:int)->None:
+        "Writes model weight histograms to Tensorboard."
+        generator, critic = self.learn.gan_trainer.generator, self.learn.gan_trainer.critic
+        self.hist_writer.write(model=generator, iteration=iteration, tbwriter=self.tbwriter, name='generator')
+        self.hist_writer.write(model=critic,    iteration=iteration, tbwriter=self.tbwriter, name='critic')
+
+    def _write_gen_model_stats(self, iteration:int)->None:
+        "Writes gradient statistics for generator to Tensorboard."
+        generator = self.learn.gan_trainer.generator
+        self.stats_writer.write(model=generator, iteration=iteration, tbwriter=self.tbwriter, name='gen_model_stats')
+        self.gen_stats_updated = True
+
+    def _write_critic_model_stats(self, iteration:int)->None:
+        "Writes gradient statistics for critic to Tensorboard."
+        critic = self.learn.gan_trainer.critic
+        self.stats_writer.write(model=critic, iteration=iteration, tbwriter=self.tbwriter, name='crit_model_stats')
+        self.crit_stats_updated = True
+
+    def _write_model_stats(self, iteration:int)->None:
+        "Writes gradient statistics to Tensorboard."
+        # We don't want to write stats when model is not iterated on and hence has zeroed out gradients
+        gen_mode = self.learn.gan_trainer.gen_mode
+        if gen_mode and not self.gen_stats_updated: self._write_gen_model_stats(iteration=iteration)
+        if not gen_mode and not self.crit_stats_updated: self._write_critic_model_stats(iteration=iteration)
+
+    def _write_training_loss(self, iteration:int, last_loss:Tensor)->None:
+        "Writes training loss to Tensorboard."
+        recorder = self.learn.gan_trainer.recorder
+        if len(recorder.losses) == 0: return
+        scalar_value = to_np((recorder.losses[-1:])[0])
+        tag = self.metrics_root + 'train_loss'
+        self.tbwriter.add_scalar(tag=tag, scalar_value=scalar_value, global_step=iteration)
+
+    def _write_images(self, iteration:int)->None:
+        "Writes model generated, original and real images to Tensorboard."
+        trainer = self.learn.gan_trainer
+        #TODO:  Switching gen_mode temporarily seems a bit hacky here.  Certainly not a good side-effect.  Is there a better way?
+        gen_mode = trainer.gen_mode
+        try:
+            trainer.switch(gen_mode=True)
+            self.img_gen_vis.write(learn=self.learn, trn_batch=self.trn_batch, val_batch=self.val_batch, 
+                                    iteration=iteration, tbwriter=self.tbwriter)
+        finally: trainer.switch(gen_mode=gen_mode)
+
+    def on_batch_end(self, iteration:int, **kwargs)->None:
+        "Callback function that writes batch end appropriate data to Tensorboard."
+        super().on_batch_end(iteration=iteration, **kwargs)
+        if iteration == 0: return
+        if iteration % self.visual_iters == 0: self._write_images(iteration=iteration)
+
+    def on_backward_end(self, iteration:int, **kwargs)->None:
+        "Callback function that writes backward end appropriate data to Tensorboard."
+        if iteration == 0: return
+        self._update_batches_if_needed()
+        #TODO:  This could perhaps be implemented as queues of requests instead but that seemed like overkill. 
+        # But I'm not the biggest fan of maintaining these boolean flags either... Review pls.
+        if iteration % self.stats_iters == 0: self.gen_stats_updated, self.crit_stats_updated = False, False
+        if not (self.gen_stats_updated and self.crit_stats_updated): self._write_model_stats(iteration=iteration)
+
+class ImageGenTensorboardWriter(LearnerTensorboardWriter):
+    "Callback for non-GAN image generating Learners that writes to Tensorboard.  Extends LearnerTensorboardWriter and adds output image writes."
+    def __init__(self, learn:Learner, base_dir:Path, name:str, loss_iters:int=25, hist_iters:int=500, stats_iters:int=100, 
+                 visual_iters:int=100):
+        super().__init__(learn=learn, base_dir=base_dir, name=name, loss_iters=loss_iters, hist_iters=hist_iters, 
+                         stats_iters=stats_iters)
+        self.visual_iters = visual_iters
+        self.img_gen_vis = ImageTBWriter()
+
+    def _write_images(self, iteration:int)->None:
+        "Writes model generated, original and real images to Tensorboard"
+        self.img_gen_vis.write(learn=self.learn, trn_batch=self.trn_batch, val_batch=self.val_batch, iteration=iteration, 
+                               tbwriter=self.tbwriter)
+
+    def on_batch_end(self, iteration:int, **kwargs)->None:
+        "Callback function that writes batch end appropriate data to Tensorboard."
+        super().on_batch_end(iteration=iteration, **kwargs)
+        if iteration == 0: return
+        if iteration % self.visual_iters == 0: 
+            self._write_images(iteration=iteration)
+
+class TBWriteRequest(ABC):
+    "A request object for Tensorboard writes.  Useful for queuing up and executing asynchronous writes."
+    def __init__(self, tbwriter: SummaryWriter, iteration:int):
+        super().__init__()
+        self.tbwriter = tbwriter
+        self.iteration = iteration
+
+    @abstractmethod
+    def write(self)->None: pass   
+
+# SummaryWriter writes tend to block quite a bit.  This gets around that and greatly boosts performance.
+# Not all tensorboard writes are using this- just the ones that take a long time.  Note that the 
+# SummaryWriter does actually use a threadsafe consumer/producer design ultimately to write to Tensorboard, 
+# so writes done outside of this async loop should be fine.
+class AsyncTBWriter():
+    "Callback for GANLearners that writes to Tensorboard.  Extends LearnerTensorboardWriter and adds output image writes."
+    def __init__(self):
+        super().__init__()
+        self.stop_request = Event()
+        self.queue = Queue()
+        self.thread = Thread(target=self._queue_processor, daemon=True)
+        self.thread.start()
+
+    def request_write(self, request: TBWriteRequest)->None:
+        "Queues up an asynchronous write request to Tensorboard."
+        if self.stop_request.isSet(): return
+        self.queue.put(request)
+
+    def _queue_processor(self)->None:
+        "Processes queued up write requests asynchronously to Tensorboard."
+        while not self.stop_request.isSet():
+            while not self.queue.empty():
+                if self.stop_request.isSet(): return
+                request = self.queue.get()
+                request.write()
+            sleep(0.2)
+
+    #Provided this to stop thread explicitly or by context management (with statement) but thread should end on its own 
+    # upon program exit, due to being a daemon.  So using this is probably unecessary.
+    def close(self)->None:
+        "Stops asynchronous request queue processing thread."
+        self.stop_request.set()
+        self.thread.join()
+
+    # Nothing to do, thread already started.  Could start thread here to enforce use of context manager 
+    # (but that sounds like a pain and a bit unweildy and unecessary for actual usage)
+    def __enter__(self): pass
+
+    def __exit__(self, exc_type, exc_value, traceback): self.close()
+
+asyncTBWriter = AsyncTBWriter() 
+
+class ModelImageSet():
+    "Convenience object that holds the original, real(target) and generated versions of a single image fed to a model."
+    @staticmethod
+    def get_list_from_model(learn:Learner, ds_type:DatasetType, batch:Tuple)->[]:
+        "Factory method to convert a batch of model images to a list of ModelImageSet."
+        image_sets = []
+        x,y = batch[0],batch[1]
+        preds=[]
+        preds = learn.pred_batch(ds_type=ds_type, batch=(x,y), reconstruct=True)  
+        for orig_px, real_px, gen in zip(x,y,preds):
+            orig, real = Image(px=orig_px), Image(px=real_px)
+            image_set = ModelImageSet(orig=orig, real=real, gen=gen)
+            image_sets.append(image_set)
+        return image_sets  
+
+    def __init__(self, orig:Image, real:Image, gen:Image): self.orig, self.real, self.gen = orig, real, gen
+
+class HistogramTBRequest(TBWriteRequest):
+    "Request object for model histogram writes to Tensorboard."
+    def __init__(self, model:nn.Module, iteration:int, tbwriter:SummaryWriter, name:str):
+        super().__init__(tbwriter=tbwriter, iteration=iteration)
+        self.params = [(name, values.clone().detach().cpu()) for (name, values) in model.named_parameters()]
+        self.name = name
+
+    def _write_histogram(self, param_name:str, values)->None:
+        "Writes single model histogram to Tensorboard."
+        tag = self.name + '/weights/' + param_name
+        self.tbwriter.add_histogram(tag=tag, values=values, global_step=self.iteration)
+
+    def write(self)->None:
+        "Writes model histograms to Tensorboard."
+        for param_name, values in self.params: self._write_histogram(param_name=param_name, values=values)
+
+#If this isn't done async then this is sloooooow
+class HistogramTBWriter():
+    "Writes model histograms to Tensorboard."
+    def __init__(self): super().__init__()
+
+    def write(self, model:nn.Module, iteration:int, tbwriter:SummaryWriter, name:str='model')->None:
+        "Writes model histograms to Tensorboard."
+        request = HistogramTBRequest(model=model, iteration=iteration, tbwriter=tbwriter, name=name)
+        asyncTBWriter.request_write(request)
+
+class ModelStatsTBRequest(TBWriteRequest):
+    "Request object for model gradient statistics writes to Tensorboard."
+    def __init__(self, model:nn.Module, iteration:int, tbwriter:SummaryWriter, name:str):
+        super().__init__(tbwriter=tbwriter, iteration=iteration)
+        self.gradients = [x.grad.clone().detach().cpu() for x in model.parameters() if x.grad is not None]
+        self.name = name
+
+    def _add_gradient_scalar(self, name:str, scalar_value)->None:
+        "Writes a single scalar value for a gradient statistic to Tensorboard."
+        tag = self.name + '/gradients/' + name
+        self.tbwriter.add_scalar(tag=tag, scalar_value=scalar_value, global_step=self.iteration)
+
+    def _write_avg_norm(self, norms:[])->None:
+        "Writes the average norm of the gradients to Tensorboard."
+        avg_norm = sum(norms)/len(self.gradients)
+        self._add_gradient_scalar('avg_norm', scalar_value=avg_norm)
+
+    def _write_median_norm(self, norms:[])->None:
+        "Writes the median norm of the gradients to Tensorboard."
+        median_norm = statistics.median(norms)
+        self._add_gradient_scalar('median_norm', scalar_value=median_norm)
+
+    def _write_max_norm(self, norms:[])->None:
+        "Writes the maximum norm of the gradients to Tensorboard."
+        max_norm = max(norms)
+        self._add_gradient_scalar('max_norm', scalar_value=max_norm)
+
+    def _write_min_norm(self, norms:[])->None:
+        "Writes the minimum norm of the gradients to Tensorboard."
+        min_norm = min(norms)
+        self._add_gradient_scalar('min_norm', scalar_value=min_norm)
+
+    def _write_num_zeros(self)->None:
+        "Writes the number of zeroes in the gradients to Tensorboard."
+        gradient_nps = [to_np(x.data) for x in self.gradients]
+        num_zeros = sum((np.asarray(x) == 0.0).sum() for x in gradient_nps)
+        self._add_gradient_scalar('num_zeros', scalar_value=num_zeros)
+
+    def _write_avg_gradient(self)->None:
+        "Writes the average of the gradients to Tensorboard."
+        avg_gradient = sum(x.data.mean() for x in self.gradients)/len(self.gradients)
+        self._add_gradient_scalar('avg_gradient', scalar_value=avg_gradient)
+
+    def _write_median_gradient(self)->None:
+        "Writes the median of the gradients to Tensorboard."
+        median_gradient = statistics.median(x.data.median() for x in self.gradients)
+        self._add_gradient_scalar('median_gradient', scalar_value=median_gradient)
+
+    def _write_max_gradient(self)->None:
+        "Writes the maximum of the gradients to Tensorboard."
+        max_gradient = max(x.data.max() for x in self.gradients)
+        self._add_gradient_scalar('max_gradient', scalar_value=max_gradient)
+
+    def _write_min_gradient(self)->None:
+        "Writes the minimum of the gradients to Tensorboard."
+        min_gradient = min(x.data.min() for x in self.gradients)
+        self._add_gradient_scalar('min_gradient', scalar_value=min_gradient)
+
+    def write(self)->None:
+        "Writes model gradient statistics to Tensorboard."
+        if len(self.gradients) == 0: return
+        norms = [x.data.norm() for x in self.gradients]
+        self._write_avg_norm(norms=norms)
+        self._write_median_norm(norms=norms)
+        self._write_max_norm(norms=norms)
+        self._write_min_norm(norms=norms)
+        self._write_num_zeros()
+        self._write_avg_gradient()
+        self._write_median_gradient()
+        self._write_max_gradient()
+        self._write_min_gradient()
+
+class ModelStatsTBWriter():
+    "Writes model gradient statistics to Tensorboard."
+    def write(self, model:nn.Module, iteration:int, tbwriter:SummaryWriter, name:str='model_stats')->None:
+        "Writes model gradient statistics to Tensorboard."
+        request = ModelStatsTBRequest(model=model, iteration=iteration, tbwriter=tbwriter, name=name)
+        asyncTBWriter.request_write(request)
+
+class ImageTBRequest(TBWriteRequest):
+    "Request object for model image output writes to Tensorboard."
+    def __init__(self, learn:Learner, batch:Tuple, iteration:int, tbwriter:SummaryWriter, ds_type:DatasetType):
+        super().__init__(tbwriter=tbwriter, iteration=iteration)
+        self.image_sets = ModelImageSet.get_list_from_model(learn=learn, batch=batch, ds_type=ds_type)
+        self.ds_type = ds_type
+
+    def _write_images(self, name:str, images:[Tensor])->None:
+        "Writes list of images as tensors to Tensorboard."
+        tag = self.ds_type.name + ' ' + name
+        self.tbwriter.add_image(tag=tag, img_tensor=vutils.make_grid(images, normalize=True), global_step=self.iteration)
+
+    def _get_image_tensors(self)->([Tensor], [Tensor], [Tensor]):
+        "Gets list of image tensors from lists of Image objects, as a tuple of original, generated and real(target) images."
+        orig_images, gen_images, real_images = [], [], []
+        for image_set in self.image_sets:
+            orig_images.append(image_set.orig.px)
+            gen_images.append(image_set.gen.px)
+            real_images.append(image_set.real.px) 
+        return orig_images, gen_images, real_images  
+
+    def write(self)->None:
+        "Writes original, generated and real(target) images to Tensorboard."
+        orig_images, gen_images, real_images = self._get_image_tensors()
+        self._write_images(name='orig images', images=orig_images)
+        self._write_images(name='gen images',  images=gen_images)
+        self._write_images(name='real images', images=real_images)
+
+#If this isn't done async then this is noticeably slower
+class ImageTBWriter():
+    "Writes model image output to Tensorboard."
+    def __init__(self): super().__init__()
+
+    def write(self, learn:Learner, trn_batch:Tuple, val_batch:Tuple, iteration:int, tbwriter:SummaryWriter)->None:
+        "Writes training and validation batch images to Tensorboard."
+        self._write_for_dstype(learn=learn, batch=val_batch, iteration=iteration, tbwriter=tbwriter, ds_type=DatasetType.Valid)
+        self._write_for_dstype(learn=learn, batch=trn_batch, iteration=iteration, tbwriter=tbwriter, ds_type=DatasetType.Train)
+
+    def _write_for_dstype(self, learn:Learner, batch:Tuple, iteration:int, tbwriter:SummaryWriter, ds_type:DatasetType)->None:
+        "Writes batch images of specified DatasetType to Tensorboard."
+        request = ImageTBRequest(learn=learn, batch=batch, iteration=iteration, tbwriter=tbwriter, ds_type=ds_type)
+        asyncTBWriter.request_write(request)
+
+class GraphTBRequest(TBWriteRequest):
+    "Request object for model histogram writes to Tensorboard."
+    def __init__(self, model:nn.Module, tbwriter:SummaryWriter, input_to_model:torch.Tensor):
+        super().__init__(tbwriter=tbwriter, iteration=0)
+        self.model,self.input_to_model = model,input_to_model
+
+    def write(self)->None:
+        "Writes single model graph to Tensorboard."
+        self.tbwriter.add_graph(model=self.model, input_to_model=self.input_to_model)
+
+class GraphTBWriter():
+    "Writes model network graph to Tensorboard."
+    def write(self, model:nn.Module, tbwriter:SummaryWriter, input_to_model:torch.Tensor)->None:
+        "Writes model graph to Tensorboard."
+        request = GraphTBRequest(model=model, tbwriter=tbwriter, input_to_model=input_to_model)
+        asyncTBWriter.request_write(request)

DeOldify/fastai/callbacks/tracker.py

@@ -0,0 +1,153 @@
+# Contribution from @fredguth, https://github.com/fredguth/fastai_playground.
+
+from fastai.torch_core import *
+from fastai.callback import *
+from fastai.basic_train import *
+
+__all__ = ['TerminateOnNaNCallback', 'EarlyStoppingCallback', 'SaveModelCallback', 'TrackerCallback',
+        'ReduceLROnPlateauCallback', 'TrackEpochCallback' ]
+
+class TerminateOnNaNCallback(Callback):
+    "A `Callback` that terminates training if loss is NaN."
+
+    def __init__(self):
+        self.stop = False
+
+    def on_batch_end(self, last_loss, epoch, num_batch, **kwargs:Any)->None:
+        "Test if `last_loss` is NaN and interrupts training."
+        if self.stop: return True #to skip validation after stopping during training
+        if torch.isnan(last_loss):
+            print (f'Epoch/Batch ({epoch}/{num_batch}): Invalid loss, terminating training.')
+            return {'stop_epoch': True, 'stop_training': True, 'skip_validate': True}
+
+class TrackerCallback(LearnerCallback):
+    "A `LearnerCallback` that keeps track of the best value in `monitor`."
+    def __init__(self, learn:Learner, monitor:str='valid_loss', mode:str='auto'):
+        super().__init__(learn)
+        self.monitor,self.mode = monitor,mode
+        if self.mode not in ['auto', 'min', 'max']:
+            warn(f'{self.__class__} mode {self.mode} is invalid, falling back to "auto" mode.')
+            self.mode = 'auto'
+        mode_dict = {'min': np.less, 'max':np.greater}
+        mode_dict['auto'] = np.less if 'loss' in self.monitor else np.greater
+        self.operator = mode_dict[self.mode]
+
+    def on_train_begin(self, **kwargs:Any)->None:
+        "Initializes the best value."
+        self.best = float('inf') if self.operator == np.less else -float('inf')
+
+    def get_monitor_value(self):
+        "Pick the monitored value."
+        if self.monitor=='trn_loss' and len(self.learn.recorder.losses) == 0: return None
+        elif len(self.learn.recorder.val_losses) == 0: return None
+        values = {'train_loss':self.learn.recorder.losses[-1].cpu().numpy(),
+                  'valid_loss':self.learn.recorder.val_losses[-1]}
+        if values['valid_loss'] is None: return
+        if self.learn.recorder.metrics:
+            for m, n in zip(self.learn.recorder.metrics[-1],self.learn.recorder.names[3:-1]):
+                values[n] = m
+        if values.get(self.monitor) is None:
+            warn(f'{self.__class__} conditioned on metric `{self.monitor}` which is not available. Available metrics are: {", ".join(map(str, self.learn.recorder.names[1:-1]))}')
+        return values.get(self.monitor)
+
+class EarlyStoppingCallback(TrackerCallback):
+    "A `TrackerCallback` that terminates training when monitored quantity stops improving."
+    def __init__(self, learn:Learner, monitor:str='valid_loss', mode:str='auto', min_delta:int=0, patience:int=0):
+        super().__init__(learn, monitor=monitor, mode=mode)
+        self.min_delta,self.patience = min_delta,patience
+        if self.operator == np.less:  self.min_delta *= -1
+
+    def on_train_begin(self, **kwargs:Any)->None:
+        "Initialize inner arguments."
+        self.wait = 0
+        super().on_train_begin(**kwargs)
+
+    def on_epoch_end(self, epoch, **kwargs:Any)->None:
+        "Compare the value monitored to its best score and maybe stop training."
+        current = self.get_monitor_value()
+        if current is None: return
+        if self.operator(current - self.min_delta, self.best):
+            self.best,self.wait = current,0
+        else:
+            self.wait += 1
+            if self.wait > self.patience:
+                print(f'Epoch {epoch}: early stopping')
+                return {"stop_training":True}
+
+class SaveModelCallback(TrackerCallback):
+    "A `TrackerCallback` that saves the model when monitored quantity is best."
+    def __init__(self, learn:Learner, monitor:str='valid_loss', mode:str='auto', every:str='improvement', name:str='bestmodel'):
+        super().__init__(learn, monitor=monitor, mode=mode)
+        self.every,self.name = every,name
+        if self.every not in ['improvement', 'epoch']:
+            warn(f'SaveModel every {self.every} is invalid, falling back to "improvement".')
+            self.every = 'improvement'
+                 
+    def jump_to_epoch(self, epoch:int)->None:
+        try: 
+            self.learn.load(f'{self.name}_{epoch-1}', purge=False)
+            print(f"Loaded {self.name}_{epoch-1}")
+        except: print(f'Model {self.name}_{epoch-1} not found.')
+
+    def on_epoch_end(self, epoch:int, **kwargs:Any)->None:
+        "Compare the value monitored to its best score and maybe save the model."
+        if self.every=="epoch": self.learn.save(f'{self.name}_{epoch}')
+        else: #every="improvement"
+            current = self.get_monitor_value()
+            if current is not None and self.operator(current, self.best):
+                print(f'Better model found at epoch {epoch} with {self.monitor} value: {current}.')
+                self.best = current
+                self.learn.save(f'{self.name}')
+
+    def on_train_end(self, **kwargs):
+        "Load the best model."
+        if self.every=="improvement" and (self.learn.path/f'{self.learn.model_dir}/{self.name}.pth').is_file():
+            self.learn.load(f'{self.name}', purge=False)
+
+class ReduceLROnPlateauCallback(TrackerCallback):
+    "A `TrackerCallback` that reduces learning rate when a metric has stopped improving."
+    def __init__(self, learn:Learner, monitor:str='valid_loss', mode:str='auto', patience:int=0, factor:float=0.2,
+                 min_delta:int=0):
+        super().__init__(learn, monitor=monitor, mode=mode)
+        self.patience,self.factor,self.min_delta = patience,factor,min_delta
+        if self.operator == np.less:  self.min_delta *= -1
+
+    def on_train_begin(self, **kwargs:Any)->None:
+        "Initialize inner arguments."
+        self.wait, self.opt = 0, self.learn.opt
+        super().on_train_begin(**kwargs)
+
+    def on_epoch_end(self, epoch, **kwargs:Any)->None:
+        "Compare the value monitored to its best and maybe reduce lr."
+        current = self.get_monitor_value()
+        if current is None: return
+        if self.operator(current - self.min_delta, self.best): self.best,self.wait = current,0
+        else:
+            self.wait += 1
+            if self.wait > self.patience:
+                self.opt.lr *= self.factor
+                self.wait = 0
+                print(f'Epoch {epoch}: reducing lr to {self.opt.lr}')
+
+
+class TrackEpochCallback(LearnerCallback):
+    _order = -20 #Need to run before fit_one_cycle
+    def __init__(self, learn:Learner, name:str='epoch', epoch_offset:int=None):
+        "Store completed epoch number in `learn.model_dir/name`."
+        super().__init__(learn)
+        learn._test_writeable_path()
+        self.path = learn.path/learn.model_dir/name
+        if epoch_offset is None:
+            if os.path.isfile(self.path):
+                 with self.path.open('r') as f:
+                     try:    self.start_epoch = int(f.read())+1
+                     except: self.start_epoch = 0
+            else: self.start_epoch = 0
+                
+    def on_train_begin(self, **kwargs:Any):
+        return {'epoch': self.start_epoch}
+
+    def on_epoch_end(self, epoch, **kwargs:Any)->None:
+        with self.path.open('w') as f: f.write(f'{epoch}')
+
+    def restart(self): os.remove(self.path)

DeOldify/fastai/collab.py

@@ -0,0 +1,106 @@
+"Module support for Collaborative Filtering"
+from .tabular import *
+from . import tabular
+
+__all__ = [*tabular.__all__, 'EmbeddingDotBias', 'EmbeddingNN', 'collab_learner', 'CollabDataBunch', 'CollabLine',
+           'CollabList', 'CollabLearner']
+
+class CollabProcessor(TabularProcessor):
+    "Subclass `TabularProcessor for `process_one`."
+    def process_one(self, item):
+        res = super().process_one(item)
+        return CollabLine(res.cats,res.conts,res.classes,res.names)
+
+class CollabLine(TabularLine):
+    "Base item for collaborative filtering, subclasses `TabularLine`."
+    def __init__(self, cats, conts, classes, names):
+        super().__init__(cats, conts, classes, names)
+        self.data = [self.data[0][0],self.data[0][1]]
+
+class CollabList(TabularList):
+    "Base `ItemList` for collaborative filtering, subclasses `TabularList`."
+    _item_cls,_label_cls,_processor = CollabLine,FloatList,CollabProcessor
+
+    def reconstruct(self, t:Tensor): return CollabLine(tensor(t), tensor([]), self.classes, self.col_names)
+
+class EmbeddingNN(TabularModel):
+    "Subclass `TabularModel` to create a NN suitable for collaborative filtering."
+    def __init__(self, emb_szs:ListSizes, layers:Collection[int]=None, ps:Collection[float]=None,
+                 emb_drop:float=0., y_range:OptRange=None, use_bn:bool=True, bn_final:bool=False):
+        super().__init__(emb_szs=emb_szs, n_cont=0, out_sz=1, layers=layers, ps=ps, emb_drop=emb_drop, y_range=y_range,
+                         use_bn=use_bn, bn_final=bn_final)
+
+    def forward(self, users:LongTensor, items:LongTensor) -> Tensor:
+        return super().forward(torch.stack([users,items], dim=1), None)
+
+class EmbeddingDotBias(Module):
+    "Base dot model for collaborative filtering."
+    def __init__(self, n_factors:int, n_users:int, n_items:int, y_range:Tuple[float,float]=None):
+        self.y_range = y_range
+        (self.u_weight, self.i_weight, self.u_bias, self.i_bias) = [embedding(*o) for o in [
+            (n_users, n_factors), (n_items, n_factors), (n_users,1), (n_items,1)
+        ]]
+
+    def forward(self, users:LongTensor, items:LongTensor) -> Tensor:
+        dot = self.u_weight(users)* self.i_weight(items)
+        res = dot.sum(1) + self.u_bias(users).squeeze() + self.i_bias(items).squeeze()
+        if self.y_range is None: return res
+        return torch.sigmoid(res) * (self.y_range[1]-self.y_range[0]) + self.y_range[0]
+
+class CollabDataBunch(DataBunch):
+    "Base `DataBunch` for collaborative filtering."
+    @classmethod
+    def from_df(cls, ratings:DataFrame, valid_pct:float=0.2, user_name:Optional[str]=None, item_name:Optional[str]=None,
+                rating_name:Optional[str]=None, test:DataFrame=None, seed:int=None, path:PathOrStr='.', bs:int=64, 
+                val_bs:int=None, num_workers:int=defaults.cpus, dl_tfms:Optional[Collection[Callable]]=None, 
+                device:torch.device=None, collate_fn:Callable=data_collate, no_check:bool=False) -> 'CollabDataBunch':
+        "Create a `DataBunch` suitable for collaborative filtering from `ratings`."
+        user_name   = ifnone(user_name,  ratings.columns[0])
+        item_name   = ifnone(item_name,  ratings.columns[1])
+        rating_name = ifnone(rating_name,ratings.columns[2])
+        cat_names = [user_name,item_name]
+        src = (CollabList.from_df(ratings, cat_names=cat_names, procs=Categorify)
+               .split_by_rand_pct(valid_pct=valid_pct, seed=seed).label_from_df(cols=rating_name))
+        if test is not None: src.add_test(CollabList.from_df(test, cat_names=cat_names))
+        return src.databunch(path=path, bs=bs, val_bs=val_bs, num_workers=num_workers, device=device, 
+                             collate_fn=collate_fn, no_check=no_check)
+
+class CollabLearner(Learner):
+    "`Learner` suitable for collaborative filtering."
+    def get_idx(self, arr:Collection, is_item:bool=True):
+        "Fetch item or user (based on `is_item`) for all in `arr`. (Set model to `cpu` and no grad.)"
+        m = self.model.eval().cpu()
+        requires_grad(m,False)
+        u_class,i_class = self.data.train_ds.x.classes.values()
+        classes = i_class if is_item else u_class
+        c2i = {v:k for k,v in enumerate(classes)}
+        try: return tensor([c2i[o] for o in arr])
+        except Exception as e: 
+            print(f"""You're trying to access {'an item' if is_item else 'a user'} that isn't in the training data.
+                  If it was in your original data, it may have been split such that it's only in the validation set now.""")
+
+    def bias(self, arr:Collection, is_item:bool=True):
+        "Bias for item or user (based on `is_item`) for all in `arr`. (Set model to `cpu` and no grad.)"
+        idx = self.get_idx(arr, is_item)
+        m = self.model
+        layer = m.i_bias if is_item else m.u_bias
+        return layer(idx).squeeze()
+
+    def weight(self, arr:Collection, is_item:bool=True):
+        "Bias for item or user (based on `is_item`) for all in `arr`. (Set model to `cpu` and no grad.)"
+        idx = self.get_idx(arr, is_item)
+        m = self.model
+        layer = m.i_weight if is_item else m.u_weight
+        return layer(idx)
+
+def collab_learner(data, n_factors:int=None, use_nn:bool=False, emb_szs:Dict[str,int]=None, layers:Collection[int]=None, 
+                   ps:Collection[float]=None, emb_drop:float=0., y_range:OptRange=None, use_bn:bool=True, 
+                   bn_final:bool=False, **learn_kwargs)->Learner:
+    "Create a Learner for collaborative filtering on `data`."
+    emb_szs = data.get_emb_szs(ifnone(emb_szs, {}))
+    u,m = data.train_ds.x.classes.values()
+    if use_nn: model = EmbeddingNN(emb_szs=emb_szs, layers=layers, ps=ps, emb_drop=emb_drop, y_range=y_range, 
+                                   use_bn=use_bn, bn_final=bn_final, **learn_kwargs)
+    else:      model = EmbeddingDotBias(n_factors, len(u), len(m), y_range=y_range)
+    return CollabLearner(data, model, **learn_kwargs)
+

DeOldify/fastai/core.py

@@ -0,0 +1,535 @@
+"`fastai.core` contains essential util functions to format and split data"
+from .imports.core import *
+
+warnings.filterwarnings("ignore", message="numpy.dtype size changed")
+warnings.filterwarnings("ignore", message="numpy.ufunc size changed")
+
+AnnealFunc = Callable[[Number, Number, float], Number]
+ArgStar = Collection[Any]
+BatchSamples = Collection[Tuple[Collection[int], int]]
+DataFrameOrChunks = Union[DataFrame, pd.io.parsers.TextFileReader]
+FilePathList = Collection[Path]
+Floats = Union[float, Collection[float]]
+ImgLabel = str
+ImgLabels = Collection[ImgLabel]
+IntsOrStrs = Union[int, Collection[int], str, Collection[str]]
+KeyFunc = Callable[[int], int]
+KWArgs = Dict[str, Any]
+ListOrItem = Union[Collection[Any], int, float, str]
+ListRules = Collection[Callable[[str], str]]
+ListSizes = Collection[Tuple[int, int]]
+NPArrayableList = Collection[Union[np.ndarray, list]]
+NPArrayList = Collection[np.ndarray]
+NPArrayMask = np.ndarray
+NPImage = np.ndarray
+OptDataFrame = Optional[DataFrame]
+OptListOrItem = Optional[ListOrItem]
+OptRange = Optional[Tuple[float, float]]
+OptStrTuple = Optional[Tuple[str, str]]
+OptStats = Optional[Tuple[np.ndarray, np.ndarray]]
+PathOrStr = Union[Path, str]
+PathLikeOrBinaryStream = Union[PathOrStr, BufferedWriter, BytesIO]
+PBar = Union[MasterBar, ProgressBar]
+Point = Tuple[float, float]
+Points = Collection[Point]
+Sizes = List[List[int]]
+SplitArrayList = List[Tuple[np.ndarray, np.ndarray]]
+StartOptEnd = Union[float, Tuple[float, float]]
+StrList = Collection[str]
+Tokens = Collection[Collection[str]]
+OptStrList = Optional[StrList]
+np.set_printoptions(precision=6, threshold=50, edgeitems=4, linewidth=120)
+
+
+def num_cpus() -> int:
+    "Get number of cpus"
+    try:
+        return len(os.sched_getaffinity(0))
+    except AttributeError:
+        return os.cpu_count()
+
+
+_default_cpus = min(16, num_cpus())
+defaults = SimpleNamespace(
+    cpus=_default_cpus, cmap='viridis', return_fig=False, silent=False)
+
+
+def is_listy(x: Any) -> bool: return isinstance(x, (tuple, list))
+def is_tuple(x: Any) -> bool: return isinstance(x, tuple)
+def is_dict(x: Any) -> bool: return isinstance(x, dict)
+def is_pathlike(x: Any) -> bool: return isinstance(x, (str, Path))
+def noop(x): return x
+
+
+class PrePostInitMeta(type):
+    "A metaclass that calls optional `__pre_init__` and `__post_init__` methods"
+    def __new__(cls, name, bases, dct):
+        x = super().__new__(cls, name, bases, dct)
+        old_init = x.__init__
+        def _pass(self): pass
+
+        @functools.wraps(old_init)
+        def _init(self, *args, **kwargs):
+            self.__pre_init__()
+            old_init(self, *args, **kwargs)
+            self.__post_init__()
+        x.__init__ = _init
+        if not hasattr(x, '__pre_init__'):
+            x.__pre_init__ = _pass
+        if not hasattr(x, '__post_init__'):
+            x.__post_init__ = _pass
+        return x
+
+
+def chunks(l: Collection, n: int) -> Iterable:
+    "Yield successive `n`-sized chunks from `l`."
+    for i in range(0, len(l), n):
+        yield l[i:i+n]
+
+
+def recurse(func: Callable, x: Any, *args, **kwargs) -> Any:
+    if is_listy(x):
+        return [recurse(func, o, *args, **kwargs) for o in x]
+    if is_dict(x):
+        return {k: recurse(func, v, *args, **kwargs) for k, v in x.items()}
+    return func(x, *args, **kwargs)
+
+
+def first_el(x: Any) -> Any:
+    "Recursively get the first element of `x`."
+    if is_listy(x):
+        return first_el(x[0])
+    if is_dict(x):
+        return first_el(x[list(x.keys())[0]])
+    return x
+
+
+def to_int(b: Any) -> Union[int, List[int]]:
+    "Recursively convert `b` to an int or list/dict of ints; raises exception if not convertible."
+    return recurse(lambda x: int(x), b)
+
+
+def ifnone(a: Any, b: Any) -> Any:
+    "`a` if `a` is not None, otherwise `b`."
+    return b if a is None else a
+
+
+def is1d(a: Collection) -> bool:
+    "Return `True` if `a` is one-dimensional"
+    return len(a.shape) == 1 if hasattr(a, 'shape') else len(np.array(a).shape) == 1
+
+
+def uniqueify(x: Series, sort: bool = False) -> List:
+    "Return sorted unique values of `x`."
+    res = list(OrderedDict.fromkeys(x).keys())
+    if sort:
+        res.sort()
+    return res
+
+
+def idx_dict(a):
+    "Create a dictionary value to index from `a`."
+    return {v: k for k, v in enumerate(a)}
+
+
+def find_classes(folder: Path) -> FilePathList:
+    "List of label subdirectories in imagenet-style `folder`."
+    classes = [d for d in folder.iterdir()
+               if d.is_dir() and not d.name.startswith('.')]
+    assert (len(classes) > 0)
+    return sorted(classes, key=lambda d: d.name)
+
+
+def arrays_split(mask: NPArrayMask, *arrs: NPArrayableList) -> SplitArrayList:
+    "Given `arrs` is [a,b,...] and `mask`index - return[(a[mask],a[~mask]),(b[mask],b[~mask]),...]."
+    assert all([len(arr) == len(arrs[0])
+               for arr in arrs]), 'All arrays should have same length'
+    mask = array(mask)
+    return list(zip(*[(a[mask], a[~mask]) for a in map(np.array, arrs)]))
+
+
+def random_split(valid_pct: float, *arrs: NPArrayableList) -> SplitArrayList:
+    "Randomly split `arrs` with `valid_pct` ratio. good for creating validation set."
+    assert (valid_pct >= 0 and valid_pct <=
+            1), 'Validation set percentage should be between 0 and 1'
+    is_train = np.random.uniform(size=(len(arrs[0]),)) > valid_pct
+    return arrays_split(is_train, *arrs)
+
+
+def listify(p: OptListOrItem = None, q: OptListOrItem = None):
+    "Make `p` listy and the same length as `q`."
+    if p is None:
+        p = []
+    elif isinstance(p, str):
+        p = [p]
+    elif not isinstance(p, Iterable):
+        p = [p]
+    # Rank 0 tensors in PyTorch are Iterable but don't have a length.
+    else:
+        try:
+            a = len(p)
+        except:
+            p = [p]
+    n = q if type(q) == int else len(p) if q is None else len(q)
+    if len(p) == 1:
+        p = p * n
+    assert len(p) == n, f'List len mismatch ({len(p)} vs {n})'
+    return list(p)
+
+
+_camel_re1 = re.compile('(.)([A-Z][a-z]+)')
+_camel_re2 = re.compile('([a-z0-9])([A-Z])')
+
+
+def camel2snake(name: str) -> str:
+    "Change `name` from camel to snake style."
+    s1 = re.sub(_camel_re1, r'\1_\2', name)
+    return re.sub(_camel_re2, r'\1_\2', s1).lower()
+
+
+def even_mults(start: float, stop: float, n: int) -> np.ndarray:
+    "Build log-stepped array from `start` to `stop` in `n` steps."
+    mult = stop/start
+    step = mult**(1/(n-1))
+    return np.array([start*(step**i) for i in range(n)])
+
+
+def extract_kwargs(names: Collection[str], kwargs: KWArgs):
+    "Extract the keys in `names` from the `kwargs`."
+    new_kwargs = {}
+    for arg_name in names:
+        if arg_name in kwargs:
+            arg_val = kwargs.pop(arg_name)
+            new_kwargs[arg_name] = arg_val
+    return new_kwargs, kwargs
+
+
+def partition(a: Collection, sz: int) -> List[Collection]:
+    "Split iterables `a` in equal parts of size `sz`"
+    return [a[i:i+sz] for i in range(0, len(a), sz)]
+
+
+def partition_by_cores(a: Collection, n_cpus: int) -> List[Collection]:
+    "Split data in `a` equally among `n_cpus` cores"
+    return partition(a, len(a)//n_cpus + 1)
+
+
+def series2cat(df: DataFrame, *col_names):
+    "Categorifies the columns `col_names` in `df`."
+    for c in listify(col_names):
+        df[c] = df[c].astype('category').cat.as_ordered()
+
+
+TfmList = Union[Callable, Collection[Callable]]
+
+
+class ItemBase():
+    "Base item type in the fastai library."
+
+    def __init__(self, data: Any): self.data = self.obj = data
+    def __repr__(self) -> str: return f'{self.__class__.__name__} {str(self)}'
+
+    def show(self, ax: plt.Axes, **kwargs):
+        "Subclass this method if you want to customize the way this `ItemBase` is shown on `ax`."
+        ax.set_title(str(self))
+
+    def apply_tfms(self, tfms: Collection, **kwargs):
+        "Subclass this method if you want to apply data augmentation with `tfms` to this `ItemBase`."
+        if tfms:
+            raise Exception(
+                f"Not implemented: you can't apply transforms to this type of item ({self.__class__.__name__})")
+        return self
+
+    def __eq__(self, other): return recurse_eq(self.data, other.data)
+
+
+def recurse_eq(arr1, arr2):
+    if is_listy(arr1):
+        return is_listy(arr2) and len(arr1) == len(arr2) and np.all([recurse_eq(x, y) for x, y in zip(arr1, arr2)])
+    else:
+        return np.all(np.atleast_1d(arr1 == arr2))
+
+
+def download_url(url: str, dest: str, overwrite: bool = False, pbar: ProgressBar = None,
+                 show_progress=True, chunk_size=1024*1024, timeout=4, retries=5) -> None:
+    "Download `url` to `dest` unless it exists and not `overwrite`."
+    if os.path.exists(dest) and not overwrite:
+        return
+
+    s = requests.Session()
+    s.mount('http://', requests.adapters.HTTPAdapter(max_retries=retries))
+    u = s.get(url, stream=True, timeout=timeout)
+    try:
+        file_size = int(u.headers["Content-Length"])
+    except:
+        show_progress = False
+
+    with open(dest, 'wb') as f:
+        nbytes = 0
+        if show_progress:
+            pbar = progress_bar(
+                range(file_size), auto_update=False, leave=False, parent=pbar)
+        try:
+            for chunk in u.iter_content(chunk_size=chunk_size):
+                nbytes += len(chunk)
+                if show_progress:
+                    pbar.update(nbytes)
+                f.write(chunk)
+        except requests.exceptions.ConnectionError as e:
+            fname = url.split('/')[-1]
+            from fastai.datasets import Config
+            data_dir = Config().data_path()
+            timeout_txt = (f'\n Download of {url} has failed after {retries} retries\n'
+                           f' Fix the download manually:\n'
+                           f'$ mkdir -p {data_dir}\n'
+                           f'$ cd {data_dir}\n'
+                           f'$ wget -c {url}\n'
+                           f'$ tar -zxvf {fname}\n\n'
+                           f'And re-run your code once the download is successful\n')
+            print(timeout_txt)
+            import sys
+            sys.exit(1)
+
+
+def range_of(x):
+    "Create a range from 0 to `len(x)`."
+    return list(range(len(x)))
+
+
+def arange_of(x):
+    "Same as `range_of` but returns an array."
+    return np.arange(len(x))
+
+
+Path.ls = lambda x: list(x.iterdir())
+
+
+def join_path(fname: PathOrStr, path: PathOrStr = '.') -> Path:
+    "Return `Path(path)/Path(fname)`, `path` defaults to current dir."
+    return Path(path)/Path(fname)
+
+
+def join_paths(fnames: FilePathList, path: PathOrStr = '.') -> Collection[Path]:
+    "Join `path` to every file name in `fnames`."
+    path = Path(path)
+    return [join_path(o, path) for o in fnames]
+
+
+def loadtxt_str(path: PathOrStr) -> np.ndarray:
+    "Return `ndarray` of `str` of lines of text from `path`."
+    with open(path, 'r') as f:
+        lines = f.readlines()
+    return np.array([l.strip() for l in lines])
+
+
+def save_texts(fname: PathOrStr, texts: Collection[str]):
+    "Save in `fname` the content of `texts`."
+    with open(fname, 'w') as f:
+        for t in texts:
+            f.write(f'{t}\n')
+
+
+def df_names_to_idx(names: IntsOrStrs, df: DataFrame):
+    "Return the column indexes of `names` in `df`."
+    if not is_listy(names):
+        names = [names]
+    if isinstance(names[0], int):
+        return names
+    return [df.columns.get_loc(c) for c in names]
+
+
+def one_hot(x: Collection[int], c: int):
+    "One-hot encode `x` with `c` classes."
+    res = np.zeros((c,), np.float32)
+    res[listify(x)] = 1.
+    return res
+
+
+def index_row(a: Union[Collection, pd.DataFrame, pd.Series], idxs: Collection[int]) -> Any:
+    "Return the slice of `a` corresponding to `idxs`."
+    if a is None:
+        return a
+    if isinstance(a, (pd.DataFrame, pd.Series)):
+        res = a.iloc[idxs]
+        if isinstance(res, (pd.DataFrame, pd.Series)):
+            return res.copy()
+        return res
+    return a[idxs]
+
+
+def func_args(func) -> bool:
+    "Return the arguments of `func`."
+    code = func.__code__
+    return code.co_varnames[:code.co_argcount]
+
+
+def has_arg(func, arg) -> bool:
+    "Check if `func` accepts `arg`."
+    return arg in func_args(func)
+
+
+def split_kwargs_by_func(kwargs, func):
+    "Split `kwargs` between those expected by `func` and the others."
+    args = func_args(func)
+    func_kwargs = {a: kwargs.pop(a) for a in args if a in kwargs}
+    return func_kwargs, kwargs
+
+
+def array(a, dtype: type = None, **kwargs) -> np.ndarray:
+    "Same as `np.array` but also handles generators. `kwargs` are passed to `np.array` with `dtype`."
+    if not isinstance(a, collections.abc.Sized) and not getattr(a, '__array_interface__', False):
+        a = list(a)
+    if np.int_ == np.int32 and dtype is None and is_listy(a) and len(a) and isinstance(a[0], int):
+        dtype = np.int64
+    return np.array(a, dtype=dtype, **kwargs)
+
+
+class EmptyLabel(ItemBase):
+    "Should be used for a dummy label."
+
+    def __init__(self): self.obj, self.data = 0, 0
+    def __str__(self): return ''
+    def __hash__(self): return hash(str(self))
+
+
+class Category(ItemBase):
+    "Basic class for single classification labels."
+
+    def __init__(self, data, obj): self.data, self.obj = data, obj
+    def __int__(self): return int(self.data)
+    def __str__(self): return str(self.obj)
+    def __hash__(self): return hash(str(self))
+
+
+class MultiCategory(ItemBase):
+    "Basic class for multi-classification labels."
+
+    def __init__(self, data, obj,
+                 raw): self.data, self.obj, self.raw = data, obj, raw
+
+    def __str__(self): return ';'.join([str(o) for o in self.obj])
+    def __hash__(self): return hash(str(self))
+
+
+class FloatItem(ItemBase):
+    "Basic class for float items."
+
+    def __init__(self, obj): self.data, self.obj = np.array(
+        obj).astype(np.float32), obj
+
+    def __str__(self): return str(self.obj)
+    def __hash__(self): return hash(str(self))
+
+
+def _treat_html(o: str) -> str:
+    o = str(o)
+    to_replace = {'\n': '\\n', '<': '&lt;', '>': '&gt;', '&': '&amp;'}
+    for k, v in to_replace.items():
+        o = o.replace(k, v)
+    return o
+
+
+def text2html_table(items: Collection[Collection[str]]) -> str:
+    "Put the texts in `items` in an HTML table, `widths` are the widths of the columns in %."
+    html_code = f"""<table border="1" class="dataframe">"""
+    html_code += f"""  <thead>\n    <tr style="text-align: right;">\n"""
+    for i in items[0]:
+        html_code += f"      <th>{_treat_html(i)}</th>"
+    html_code += f"    </tr>\n  </thead>\n  <tbody>"
+    html_code += "  <tbody>"
+    for line in items[1:]:
+        html_code += "    <tr>"
+        for i in line:
+            html_code += f"      <td>{_treat_html(i)}</td>"
+        html_code += "    </tr>"
+    html_code += "  </tbody>\n</table>"
+    return html_code
+
+
+def parallel(func, arr: Collection, max_workers: int = None, leave=False):
+    "Call `func` on every element of `arr` in parallel using `max_workers`."
+    max_workers = ifnone(max_workers, defaults.cpus)
+    if max_workers < 2:
+        results = [func(o, i) for i, o in progress_bar(
+            enumerate(arr), total=len(arr), leave=leave)]
+    else:
+        with ProcessPoolExecutor(max_workers=max_workers) as ex:
+            futures = [ex.submit(func, o, i) for i, o in enumerate(arr)]
+            results = []
+            for f in progress_bar(concurrent.futures.as_completed(futures), total=len(arr), leave=leave):
+                results.append(f.result())
+    if any([o is not None for o in results]):
+        return results
+
+
+def subplots(rows: int, cols: int, imgsize: int = 4, figsize: Optional[Tuple[int, int]] = None, title=None, **kwargs):
+    "Like `plt.subplots` but with consistent axs shape, `kwargs` passed to `fig.suptitle` with `title`"
+    figsize = ifnone(figsize, (imgsize*cols, imgsize*rows))
+    fig, axs = plt.subplots(rows, cols, figsize=figsize)
+    if rows == cols == 1:
+        axs = [[axs]]  # subplots(1,1) returns Axes, not [Axes]
+    elif (rows == 1 and cols != 1) or (cols == 1 and rows != 1):
+        axs = [axs]
+    if title is not None:
+        fig.suptitle(title, **kwargs)
+    return array(axs)
+
+
+def show_some(items: Collection, n_max: int = 5, sep: str = ','):
+    "Return the representation of the first  `n_max` elements in `items`."
+    if items is None or len(items) == 0:
+        return ''
+    res = sep.join([f'{o}' for o in items[:n_max]])
+    if len(items) > n_max:
+        res += '...'
+    return res
+
+
+def get_tmp_file(dir=None):
+    "Create and return a tmp filename, optionally at a specific path. `os.remove` when done with it."
+    with tempfile.NamedTemporaryFile(delete=False, dir=dir) as f:
+        return f.name
+
+
+def compose(funcs: List[Callable]) -> Callable:
+    "Compose `funcs`"
+    def compose_(funcs, x, *args, **kwargs):
+        for f in listify(funcs):
+            x = f(x, *args, **kwargs)
+        return x
+    return partial(compose_, funcs)
+
+
+class PrettyString(str):
+    "Little hack to get strings to show properly in Jupyter."
+
+    def __repr__(self): return self
+
+
+def float_or_x(x):
+    "Tries to convert to float, returns x if it can't"
+    try:
+        return float(x)
+    except:
+        return x
+
+
+def bunzip(fn: PathOrStr):
+    "bunzip `fn`, raising exception if output already exists"
+    fn = Path(fn)
+    assert fn.exists(), f"{fn} doesn't exist"
+    out_fn = fn.with_suffix('')
+    assert not out_fn.exists(), f"{out_fn} already exists"
+    with bz2.BZ2File(fn, 'rb') as src, out_fn.open('wb') as dst:
+        for d in iter(lambda: src.read(1024*1024), b''): dst.write(d)
+
+
+@contextmanager
+def working_directory(path: PathOrStr):
+    "Change working directory to `path` and return to previous on exit."
+    prev_cwd = Path.cwd()
+    os.chdir(path)
+    try:
+        yield
+    finally:
+        os.chdir(prev_cwd)

DeOldify/fastai/data_block.py

@@ -0,0 +1,804 @@
+from .torch_core import *
+from .basic_data import *
+from .layers import *
+from numbers import Integral
+
+__all__ = ['ItemList', 'CategoryList', 'MultiCategoryList', 'MultiCategoryProcessor', 'LabelList', 'ItemLists', 'get_files',
+           'PreProcessor', 'LabelLists', 'FloatList', 'CategoryProcessor', 'EmptyLabelList', 'MixedItem', 'MixedProcessor',
+           'MixedItemList']
+
+def _decode(df):
+    return np.array([[df.columns[i] for i,t in enumerate(x) if t==1] for x in df.values], dtype=np.object)
+
+def _maybe_squeeze(arr): return (arr if is1d(arr) else np.squeeze(arr))
+
+def _path_to_same_str(p_fn):
+    "path -> str, but same on nt+posix, for alpha-sort only"
+    s_fn = str(p_fn)
+    s_fn = s_fn.replace('\\','.')
+    s_fn = s_fn.replace('/','.')
+    return s_fn
+
+def _get_files(parent, p, f, extensions):
+    p = Path(p)#.relative_to(parent)
+    if isinstance(extensions,str): extensions = [extensions]
+    low_extensions = [e.lower() for e in extensions] if extensions is not None else None
+    res = [p/o for o in f if not o.startswith('.')
+           and (extensions is None or f'.{o.split(".")[-1].lower()}' in low_extensions)]
+    return res
+
+def get_files(path:PathOrStr, extensions:Collection[str]=None, recurse:bool=False,
+              include:Optional[Collection[str]]=None, presort:bool=False)->FilePathList:
+    "Return list of files in `path` that have a suffix in `extensions`; optionally `recurse`."
+    if recurse:
+        res = []
+        for i,(p,d,f) in enumerate(os.walk(path)):
+            # skip hidden dirs
+            if include is not None and i==0:  d[:] = [o for o in d if o in include]
+            else:                             d[:] = [o for o in d if not o.startswith('.')]
+            res += _get_files(path, p, f, extensions)
+        if presort: res = sorted(res, key=lambda p: _path_to_same_str(p), reverse=False)
+        return res
+    else:
+        f = [o.name for o in os.scandir(path) if o.is_file()]
+        res = _get_files(path, path, f, extensions)
+        if presort: res = sorted(res, key=lambda p: _path_to_same_str(p), reverse=False)
+        return res
+
+class PreProcessor():
+    "Basic class for a processor that will be applied to items at the end of the data block API."
+    def __init__(self, ds:Collection=None):  self.ref_ds = ds
+    def process_one(self, item:Any):         return item
+    def process(self, ds:Collection):        ds.items = array([self.process_one(item) for item in ds.items])
+
+PreProcessors = Union[PreProcessor, Collection[PreProcessor]]
+fastai_types[PreProcessors] = 'PreProcessors'
+
+class ItemList():
+    "A collection of items with `__len__` and `__getitem__` with `ndarray` indexing semantics."
+    _bunch,_processor,_label_cls,_square_show,_square_show_res = DataBunch,None,None,False,False
+
+    def __init__(self, items:Iterator, path:PathOrStr='.', label_cls:Callable=None, inner_df:Any=None,
+                 processor:PreProcessors=None, x:'ItemList'=None, ignore_empty:bool=False):
+        self.path = Path(path)
+        self.num_parts = len(self.path.parts)
+        self.items,self.x,self.ignore_empty = items,x,ignore_empty
+        if not isinstance(self.items,np.ndarray): self.items = array(self.items, dtype=object)
+        self.label_cls,self.inner_df,self.processor = ifnone(label_cls,self._label_cls),inner_df,processor
+        self._label_list,self._split = LabelList,ItemLists
+        self.copy_new = ['x', 'label_cls', 'path']
+
+    def __len__(self)->int: return len(self.items) or 1
+    def get(self, i)->Any:
+        "Subclass if you want to customize how to create item `i` from `self.items`."
+        return self.items[i]
+    def __repr__(self)->str:
+        items = [self[i] for i in range(min(5,len(self.items)))]
+        return f'{self.__class__.__name__} ({len(self.items)} items)\n{show_some(items)}\nPath: {self.path}'
+
+    def process(self, processor:PreProcessors=None):
+        "Apply `processor` or `self.processor` to `self`."
+        if processor is not None: self.processor = processor
+        self.processor = listify(self.processor)
+        for p in self.processor: p.process(self)
+        return self
+
+    def process_one(self, item:ItemBase, processor:PreProcessors=None):
+        "Apply `processor` or `self.processor` to `item`."
+        if processor is not None: self.processor = processor
+        self.processor = listify(self.processor)
+        for p in self.processor: item = p.process_one(item)
+        return item
+
+    def analyze_pred(self, pred:Tensor):
+        "Called on `pred` before `reconstruct` for additional preprocessing."
+        return pred
+
+    def reconstruct(self, t:Tensor, x:Tensor=None):
+        "Reconstruct one of the underlying item for its data `t`."
+        return self[0].reconstruct(t,x) if has_arg(self[0].reconstruct, 'x') else self[0].reconstruct(t)
+
+    def new(self, items:Iterator, processor:PreProcessors=None, **kwargs)->'ItemList':
+        "Create a new `ItemList` from `items`, keeping the same attributes."
+        processor = ifnone(processor, self.processor)
+        copy_d = {o:getattr(self,o) for o in self.copy_new}
+        kwargs = {**copy_d, **kwargs}
+        return self.__class__(items=items, processor=processor, **kwargs)
+
+    def add(self, items:'ItemList'):
+        self.items = np.concatenate([self.items, items.items], 0)
+        if self.inner_df is not None and items.inner_df is not None:
+            self.inner_df = pd.concat([self.inner_df, items.inner_df])
+        else: self.inner_df = self.inner_df or items.inner_df
+        return self
+
+    def __getitem__(self,idxs:int)->Any:
+        "returns a single item based if `idxs` is an integer or a new `ItemList` object if `idxs` is a range."
+        idxs = try_int(idxs)
+        if isinstance(idxs, Integral): return self.get(idxs)
+        else: return self.new(self.items[idxs], inner_df=index_row(self.inner_df, idxs))
+
+    @classmethod
+    def from_folder(cls, path:PathOrStr, extensions:Collection[str]=None, recurse:bool=True,
+                    include:Optional[Collection[str]]=None, processor:PreProcessors=None, presort:Optional[bool]=False, **kwargs)->'ItemList':
+        """Create an `ItemList` in `path` from the filenames that have a suffix in `extensions`.
+        `recurse` determines if we search subfolders."""
+        path = Path(path)
+        return cls(get_files(path, extensions, recurse=recurse, include=include, presort=presort), path=path, processor=processor, **kwargs)
+
+    @classmethod
+    def from_df(cls, df:DataFrame, path:PathOrStr='.', cols:IntsOrStrs=0, processor:PreProcessors=None, **kwargs)->'ItemList':
+        "Create an `ItemList` in `path` from the inputs in the `cols` of `df`."
+        inputs = df.iloc[:,df_names_to_idx(cols, df)]
+        assert not inputs.isna().any().any(), f"You have NaN values in column(s) {cols} of your dataframe, please fix it."
+        res = cls(items=_maybe_squeeze(inputs.values), path=path, inner_df=df, processor=processor, **kwargs)
+        return res
+
+    @classmethod
+    def from_csv(cls, path:PathOrStr, csv_name:str, cols:IntsOrStrs=0, delimiter:str=None, header:str='infer',
+                 processor:PreProcessors=None, **kwargs)->'ItemList':
+        """Create an `ItemList` in `path` from the inputs in the `cols` of `path/csv_name`"""
+        df = pd.read_csv(Path(path)/csv_name, delimiter=delimiter, header=header)
+        return cls.from_df(df, path=path, cols=cols, processor=processor, **kwargs)
+
+    def _relative_item_path(self, i): return self.items[i].relative_to(self.path)
+    def _relative_item_paths(self):   return [self._relative_item_path(i) for i in range_of(self.items)]
+
+    def use_partial_data(self, sample_pct:float=0.01, seed:int=None)->'ItemList':
+        "Use only a sample of `sample_pct`of the full dataset and an optional `seed`."
+        if seed is not None: np.random.seed(seed)
+        rand_idx = np.random.permutation(range_of(self))
+        cut = int(sample_pct * len(self))
+        return self[rand_idx[:cut]]
+
+    def to_text(self, fn:str):
+        "Save `self.items` to `fn` in `self.path`."
+        with open(self.path/fn, 'w') as f: f.writelines([f'{o}\n' for o in self._relative_item_paths()])
+
+    def filter_by_func(self, func:Callable)->'ItemList':
+        "Only keep elements for which `func` returns `True`."
+        self.items = array([o for o in self.items if func(o)])
+        return self
+
+    def filter_by_folder(self, include=None, exclude=None):
+        "Only keep filenames in `include` folder or reject the ones in `exclude`."
+        include,exclude = listify(include),listify(exclude)
+        def _inner(o):
+            if isinstance(o, Path): n = o.relative_to(self.path).parts[0]
+            else: n = o.split(os.path.sep)[len(str(self.path).split(os.path.sep))]
+            if include and not n in include: return False
+            if exclude and     n in exclude: return False
+            return True
+        return self.filter_by_func(_inner)
+
+    def filter_by_rand(self, p:float, seed:int=None):
+        "Keep random sample of `items` with probability `p` and an optional `seed`."
+        if seed is not None: set_all_seed(seed)
+        return self.filter_by_func(lambda o: rand_bool(p))
+
+    def no_split(self):
+        warn("`no_split` is deprecated, please use `split_none`.")
+        return self.split_none()
+
+    def split_none(self):
+        "Don't split the data and create an empty validation set."
+        val = self[[]]
+        val.ignore_empty = True
+        return self._split(self.path, self, val)
+
+    def split_by_list(self, train, valid):
+        "Split the data between `train` and `valid`."
+        return self._split(self.path, train, valid)
+
+    def split_by_idxs(self, train_idx, valid_idx):
+        "Split the data between `train_idx` and `valid_idx`."
+        return self.split_by_list(self[train_idx], self[valid_idx])
+
+    def split_by_idx(self, valid_idx:Collection[int])->'ItemLists':
+        "Split the data according to the indexes in `valid_idx`."
+        #train_idx = [i for i in range_of(self.items) if i not in valid_idx]
+        train_idx = np.setdiff1d(arange_of(self.items), valid_idx)
+        return self.split_by_idxs(train_idx, valid_idx)
+
+    def _get_by_folder(self, name):
+        return [i for i in range_of(self) if (self.items[i].parts[self.num_parts] if isinstance(self.items[i], Path)
+                else self.items[i].split(os.path.sep)[0]) == name ]
+
+    def split_by_folder(self, train:str='train', valid:str='valid')->'ItemLists':
+        "Split the data depending on the folder (`train` or `valid`) in which the filenames are."
+        return self.split_by_idxs(self._get_by_folder(train), self._get_by_folder(valid))
+
+    def random_split_by_pct(self, valid_pct:float=0.2, seed:int=None):
+        warn("`random_split_by_pct` is deprecated, please use `split_by_rand_pct`.")
+        return self.split_by_rand_pct(valid_pct=valid_pct, seed=seed)
+
+    def split_by_rand_pct(self, valid_pct:float=0.2, seed:int=None)->'ItemLists':
+        "Split the items randomly by putting `valid_pct` in the validation set, optional `seed` can be passed."
+        if valid_pct==0.: return self.split_none()
+        if seed is not None: np.random.seed(seed)
+        rand_idx = np.random.permutation(range_of(self))
+        cut = int(valid_pct * len(self))
+        return self.split_by_idx(rand_idx[:cut])
+
+    def split_subsets(self, train_size:float, valid_size:float, seed=None) -> 'ItemLists':
+        "Split the items into train set with size `train_size * n` and valid set with size `valid_size * n`."
+        assert 0 < train_size < 1
+        assert 0 < valid_size < 1
+        assert train_size + valid_size <= 1.
+        if seed is not None: np.random.seed(seed)
+        n = len(self.items)
+        rand_idx = np.random.permutation(range(n))
+        train_cut, valid_cut = int(train_size * n), int(valid_size * n)
+        return self.split_by_idxs(rand_idx[:train_cut], rand_idx[-valid_cut:])
+
+    def split_by_valid_func(self, func:Callable)->'ItemLists':
+        "Split the data by result of `func` (which returns `True` for validation set)."
+        valid_idx = [i for i,o in enumerate(self.items) if func(o)]
+        return self.split_by_idx(valid_idx)
+
+    def split_by_files(self, valid_names:'ItemList')->'ItemLists':
+        "Split the data by using the names in `valid_names` for validation."
+        if isinstance(self.items[0], Path): return self.split_by_valid_func(lambda o: o.name in valid_names)
+        else: return self.split_by_valid_func(lambda o: os.path.basename(o) in valid_names)
+
+    def split_by_fname_file(self, fname:PathOrStr, path:PathOrStr=None)->'ItemLists':
+        "Split the data by using the names in `fname` for the validation set. `path` will override `self.path`."
+        path = Path(ifnone(path, self.path))
+        valid_names = loadtxt_str(path/fname)
+        return self.split_by_files(valid_names)
+
+    def split_from_df(self, col:IntsOrStrs=2):
+        "Split the data from the `col` in the dataframe in `self.inner_df`."
+        valid_idx = np.where(self.inner_df.iloc[:,df_names_to_idx(col, self.inner_df)])[0]
+        return self.split_by_idx(valid_idx)
+
+    def get_label_cls(self, labels, label_cls:Callable=None, label_delim:str=None, **kwargs):
+        "Return `label_cls` or guess one from the first element of `labels`."
+        if label_cls is not None:               return label_cls
+        if self.label_cls is not None:          return self.label_cls
+        if label_delim is not None:             return MultiCategoryList
+        it = index_row(labels,0)
+        if isinstance(it, (float, np.float32)): return FloatList
+        if isinstance(try_int(it), (str, Integral)):  return CategoryList
+        if isinstance(it, Collection):          return MultiCategoryList
+        return ItemList #self.__class__
+
+    def _label_from_list(self, labels:Iterator, label_cls:Callable=None, from_item_lists:bool=False, **kwargs)->'LabelList':
+        "Label `self.items` with `labels`."
+        if not from_item_lists:
+            raise Exception("Your data isn't split, if you don't want a validation set, please use `split_none`.")
+        labels = array(labels, dtype=object)
+        label_cls = self.get_label_cls(labels, label_cls=label_cls, **kwargs)
+        y = label_cls(labels, path=self.path, **kwargs)
+        res = self._label_list(x=self, y=y)
+        return res
+
+    def label_from_df(self, cols:IntsOrStrs=1, label_cls:Callable=None, **kwargs):
+        "Label `self.items` from the values in `cols` in `self.inner_df`."
+        labels = self.inner_df.iloc[:,df_names_to_idx(cols, self.inner_df)]
+        assert labels.isna().sum().sum() == 0, f"You have NaN values in column(s) {cols} of your dataframe, please fix it."
+        if is_listy(cols) and len(cols) > 1 and (label_cls is None or label_cls == MultiCategoryList):
+            new_kwargs,label_cls = dict(one_hot=True, classes= cols),MultiCategoryList
+            kwargs = {**new_kwargs, **kwargs}
+        return self._label_from_list(_maybe_squeeze(labels), label_cls=label_cls, **kwargs)
+
+    def label_const(self, const:Any=0, label_cls:Callable=None, **kwargs)->'LabelList':
+        "Label every item with `const`."
+        return self.label_from_func(func=lambda o: const, label_cls=label_cls, **kwargs)
+
+    def label_empty(self, **kwargs):
+        "Label every item with an `EmptyLabel`."
+        kwargs['label_cls'] = EmptyLabelList
+        return self.label_from_func(func=lambda o: 0., **kwargs)
+
+    def label_from_func(self, func:Callable, label_cls:Callable=None, **kwargs)->'LabelList':
+        "Apply `func` to every input to get its label."
+        return self._label_from_list([func(o) for o in self.items], label_cls=label_cls, **kwargs)
+
+    def label_from_folder(self, label_cls:Callable=None, **kwargs)->'LabelList':
+        "Give a label to each filename depending on its folder."
+        return self.label_from_func(func=lambda o: (o.parts if isinstance(o, Path) else o.split(os.path.sep))[-2],
+                                    label_cls=label_cls, **kwargs)
+
+    def label_from_re(self, pat:str, full_path:bool=False, label_cls:Callable=None, **kwargs)->'LabelList':
+        "Apply the re in `pat` to determine the label of every filename.  If `full_path`, search in the full name."
+        pat = re.compile(pat)
+        def _inner(o):
+            s = str((os.path.join(self.path,o) if full_path else o).as_posix())
+            res = pat.search(s)
+            assert res,f'Failed to find "{pat}" in "{s}"'
+            return res.group(1)
+        return self.label_from_func(_inner, label_cls=label_cls, **kwargs)
+
+    def databunch(self, **kwargs):
+        "To throw a clear error message when the data wasn't split and labeled."
+        raise Exception("Your data is neither split nor labeled, can't turn it into a `DataBunch` yet.")
+
+class EmptyLabelList(ItemList):
+    "Basic `ItemList` for dummy labels."
+    def get(self, i): return EmptyLabel()
+    def reconstruct(self, t:Tensor, x:Tensor=None):
+        if len(t.size()) == 0: return EmptyLabel()
+        return self.x.reconstruct(t,x) if has_arg(self.x.reconstruct, 'x') else self.x.reconstruct(t)
+
+class CategoryProcessor(PreProcessor):
+    "`PreProcessor` that create `classes` from `ds.items` and handle the mapping."
+    def __init__(self, ds:ItemList):
+        self.create_classes(ds.classes)
+        self.state_attrs,self.warns = ['classes'],[]
+
+    def create_classes(self, classes):
+        self.classes = classes
+        if classes is not None: self.c2i = {v:k for k,v in enumerate(classes)}
+
+    def generate_classes(self, items):
+        "Generate classes from `items` by taking the sorted unique values."
+        return uniqueify(items, sort=True)
+
+    def process_one(self,item):
+        if isinstance(item, EmptyLabel): return item
+        res = self.c2i.get(item,None)
+        if res is None: self.warns.append(str(item))
+        return res
+
+    def process(self, ds):
+        if self.classes is None: self.create_classes(self.generate_classes(ds.items))
+        ds.classes = self.classes
+        ds.c2i = self.c2i
+        super().process(ds)
+
+    def __getstate__(self): return {n:getattr(self,n) for n in self.state_attrs}
+    def __setstate__(self, state:dict):
+        self.create_classes(state['classes'])
+        self.state_attrs = state.keys()
+        for n in state.keys():
+            if n!='classes': setattr(self, n, state[n])
+
+class CategoryListBase(ItemList):
+    "Basic `ItemList` for classification."
+    def __init__(self, items:Iterator, classes:Collection=None, **kwargs):
+        self.classes=classes
+        self.filter_missing_y = True
+        super().__init__(items, **kwargs)
+        self.copy_new.append('classes')
+
+    @property
+    def c(self): return len(self.classes)
+
+class CategoryList(CategoryListBase):
+    "Basic `ItemList` for single classification labels."
+    _processor=CategoryProcessor
+    def __init__(self, items:Iterator, classes:Collection=None, label_delim:str=None, **kwargs):
+        super().__init__(items, classes=classes, **kwargs)
+        self.loss_func = CrossEntropyFlat()
+
+    def get(self, i):
+        o = self.items[i]
+        if o is None: return None
+        return Category(o, self.classes[o])
+
+    def analyze_pred(self, pred, thresh:float=0.5): return pred.argmax()
+
+    def reconstruct(self, t):
+        return Category(t, self.classes[t])
+
+class MultiCategoryProcessor(CategoryProcessor):
+    "`PreProcessor` that create `classes` from `ds.items` and handle the mapping."
+    def __init__(self, ds:ItemList, one_hot:bool=False):
+        super().__init__(ds)
+        self.one_hot = one_hot
+        self.state_attrs.append('one_hot')
+
+    def process_one(self,item):
+        if self.one_hot or isinstance(item, EmptyLabel): return item
+        res = [super(MultiCategoryProcessor, self).process_one(o) for o in item]
+        return [r for r in res if r is not None]
+
+    def generate_classes(self, items):
+        "Generate classes from `items` by taking the sorted unique values."
+        classes = set()
+        for c in items: classes = classes.union(set(c))
+        classes = list(classes)
+        classes.sort()
+        return classes
+
+class MultiCategoryList(CategoryListBase):
+    "Basic `ItemList` for multi-classification labels."
+    _processor=MultiCategoryProcessor
+    def __init__(self, items:Iterator, classes:Collection=None, label_delim:str=None, one_hot:bool=False, **kwargs):
+        if label_delim is not None: items = array(csv.reader(items.astype(str), delimiter=label_delim))
+        super().__init__(items, classes=classes, **kwargs)
+        if one_hot:
+            assert classes is not None, "Please provide class names with `classes=...`"
+            self.processor = [MultiCategoryProcessor(self, one_hot=True)]
+        self.loss_func = BCEWithLogitsFlat()
+        self.one_hot = one_hot
+        self.copy_new += ['one_hot']
+
+    def get(self, i):
+        o = self.items[i]
+        if o is None: return None
+        if self.one_hot: return self.reconstruct(o.astype(np.float32))
+        return MultiCategory(one_hot(o, self.c), [self.classes[p] for p in o], o)
+
+    def analyze_pred(self, pred, thresh:float=0.5):
+        return (pred >= thresh).float()
+
+    def reconstruct(self, t):
+        o = [i for i in range(self.c) if t[i] == 1.]
+        return MultiCategory(t, [self.classes[p] for p in o], o)
+
+class FloatList(ItemList):
+    "`ItemList` suitable for storing the floats in items for regression. Will add a `log` if this flag is `True`."
+    def __init__(self, items:Iterator, log:bool=False, classes:Collection=None, **kwargs):
+        super().__init__(np.array(items, dtype=np.float32), **kwargs)
+        self.log = log
+        self.copy_new.append('log')
+        self.c = self.items.shape[1] if len(self.items.shape) > 1 else 1
+        self.loss_func = MSELossFlat()
+
+    def get(self, i):
+        o = super().get(i)
+        return FloatItem(np.log(o) if self.log else o)
+
+    def reconstruct(self,t): return FloatItem(t.numpy())
+
+class ItemLists():
+    "An `ItemList` for each of `train` and `valid` (optional `test`)."
+    def __init__(self, path:PathOrStr, train:ItemList, valid:ItemList):
+        self.path,self.train,self.valid,self.test = Path(path),train,valid,None
+        if not self.train.ignore_empty and len(self.train.items) == 0:
+            warn("Your training set is empty. If this is by design, pass `ignore_empty=True` to remove this warning.")
+        if not self.valid.ignore_empty and len(self.valid.items) == 0:
+            warn("""Your validation set is empty. If this is by design, use `split_none()`
+                 or pass `ignore_empty=True` when labelling to remove this warning.""")
+        if isinstance(self.train, LabelList): self.__class__ = LabelLists
+
+    def __dir__(self)->List[str]:
+        default_dir = dir(type(self)) + list(self.__dict__.keys())
+        add_ons = ['label_const', 'label_empty', 'label_from_df', 'label_from_folder', 'label_from_func',
+                   'label_from_list', 'label_from_re']
+        return default_dir + add_ons
+
+    def __repr__(self)->str:
+        return f'{self.__class__.__name__};\n\nTrain: {self.train};\n\nValid: {self.valid};\n\nTest: {self.test}'
+
+    def __getattr__(self, k):
+        ft = getattr(self.train, k)
+        if not isinstance(ft, Callable): return ft
+        fv = getattr(self.valid, k)
+        assert isinstance(fv, Callable)
+        def _inner(*args, **kwargs):
+            self.train = ft(*args, from_item_lists=True, **kwargs)
+            assert isinstance(self.train, LabelList)
+            kwargs['label_cls'] = self.train.y.__class__
+            self.valid = fv(*args, from_item_lists=True, **kwargs)
+            self.__class__ = LabelLists
+            self.process()
+            return self
+        return _inner
+
+    def __setstate__(self,data:Any): self.__dict__.update(data)
+
+    @property
+    def lists(self):
+        res = [self.train,self.valid]
+        if self.test is not None: res.append(self.test)
+        return res
+
+    def label_from_lists(self, train_labels:Iterator, valid_labels:Iterator, label_cls:Callable=None, **kwargs)->'LabelList':
+        "Use the labels in `train_labels` and `valid_labels` to label the data. `label_cls` will overwrite the default."
+        label_cls = self.train.get_label_cls(train_labels, label_cls)
+        self.train = self.train._label_list(x=self.train, y=label_cls(train_labels, **kwargs))
+        self.valid = self.valid._label_list(x=self.valid, y=self.train.y.new(valid_labels, **kwargs))
+        self.__class__ = LabelLists
+        self.process()
+        return self
+
+    def transform(self, tfms:Optional[Tuple[TfmList,TfmList]]=(None,None), **kwargs):
+        "Set `tfms` to be applied to the xs of the train and validation set."
+        if not tfms: tfms=(None,None)
+        assert is_listy(tfms) and len(tfms) == 2, "Please pass a list of two lists of transforms (train and valid)."
+        self.train.transform(tfms[0], **kwargs)
+        self.valid.transform(tfms[1], **kwargs)
+        if self.test: self.test.transform(tfms[1], **kwargs)
+        return self
+
+    def transform_y(self, tfms:Optional[Tuple[TfmList,TfmList]]=(None,None), **kwargs):
+        "Set `tfms` to be applied to the ys of the train and validation set."
+        if not tfms: tfms=(None,None)
+        self.train.transform_y(tfms[0], **kwargs)
+        self.valid.transform_y(tfms[1], **kwargs)
+        if self.test: self.test.transform_y(tfms[1], **kwargs)
+        return self
+
+    def databunch(self, **kwargs):
+        "To throw a clear error message when the data wasn't labeled."
+        raise Exception("Your data isn't labeled, can't turn it into a `DataBunch` yet!")
+
+class LabelLists(ItemLists):
+    "A `LabelList` for each of `train` and `valid` (optional `test`)."
+    def get_processors(self):
+        "Read the default class processors if none have been set."
+        procs_x,procs_y = listify(self.train.x._processor),listify(self.train.y._processor)
+        xp = ifnone(self.train.x.processor, [p(ds=self.train.x) for p in procs_x])
+        yp = ifnone(self.train.y.processor, [p(ds=self.train.y) for p in procs_y])
+        return xp,yp
+
+    def process(self):
+        "Process the inner datasets."
+        xp,yp = self.get_processors()
+        for ds,n in zip(self.lists, ['train','valid','test']): ds.process(xp, yp, name=n)
+        #progress_bar clear the outputs so in some case warnings issued during processing disappear.
+        for ds in self.lists:
+            if getattr(ds, 'warn', False): warn(ds.warn)
+        return self
+
+    def filter_by_func(self, func:Callable):
+        for ds in self.lists: ds.filter_by_func(func)
+        return self
+
+    def databunch(self, path:PathOrStr=None, bs:int=64, val_bs:int=None, num_workers:int=defaults.cpus,
+                  dl_tfms:Optional[Collection[Callable]]=None, device:torch.device=None, collate_fn:Callable=data_collate,
+                  no_check:bool=False, **kwargs)->'DataBunch':
+        "Create an `DataBunch` from self, `path` will override `self.path`, `kwargs` are passed to `DataBunch.create`."
+        path = Path(ifnone(path, self.path))
+        data = self.x._bunch.create(self.train, self.valid, test_ds=self.test, path=path, bs=bs, val_bs=val_bs,
+                                    num_workers=num_workers, dl_tfms=dl_tfms, device=device, collate_fn=collate_fn, no_check=no_check, **kwargs)
+        if getattr(self, 'normalize', False):#In case a normalization was serialized
+            norm = self.normalize
+            data.normalize((norm['mean'], norm['std']), do_x=norm['do_x'], do_y=norm['do_y'])
+        data.label_list = self
+        return data
+
+    def add_test(self, items:Iterator, label:Any=None, tfms=None, tfm_y=None):
+        "Add test set containing `items` with an arbitrary `label`."
+        # if no label passed, use label of first training item
+        if label is None: labels = EmptyLabelList([0] * len(items))
+        else: labels = self.valid.y.new([label] * len(items)).process()
+        if isinstance(items, MixedItemList): items = self.valid.x.new(items.item_lists, inner_df=items.inner_df).process()
+        elif isinstance(items, ItemList): items = self.valid.x.new(items.items, inner_df=items.inner_df).process()
+        else: items = self.valid.x.new(items).process()
+        self.test = self.valid.new(items, labels, tfms=tfms, tfm_y=tfm_y)
+        return self
+
+    def add_test_folder(self, test_folder:str='test', label:Any=None, tfms=None, tfm_y=None):
+        "Add test set containing items from `test_folder` and an arbitrary `label`."
+        # note: labels will be ignored if available in the test dataset
+        items = self.x.__class__.from_folder(self.path/test_folder)
+        return self.add_test(items.items, label=label, tfms=tfms, tfm_y=tfm_y)
+
+    @classmethod
+    def load_state(cls, path:PathOrStr, state:dict):
+        "Create a `LabelLists` with empty sets from the serialized `state`."
+        path = Path(path)
+        train_ds = LabelList.load_state(path, state)
+        valid_ds = LabelList.load_state(path, state)
+        return LabelLists(path, train=train_ds, valid=valid_ds)
+
+    @classmethod
+    def load_empty(cls, path:PathOrStr, fn:PathOrStr='export.pkl'):
+        "Create a `LabelLists` with empty sets from the serialized file in `path/fn`."
+        path = Path(path)
+        state = torch.load(open(path/fn, 'rb'))
+        return LabelLists.load_state(path, state)
+
+def _check_kwargs(ds:ItemList, tfms:TfmList, **kwargs):
+    tfms = listify(tfms)
+    if (tfms is None or len(tfms) == 0) and len(kwargs) == 0: return
+    if len(ds.items) >= 1:
+        x = ds[0]
+        try: x.apply_tfms(tfms, **kwargs)
+        except Exception as e:
+            raise Exception(f"It's not possible to apply those transforms to your dataset:\n {e}")
+
+class LabelList(Dataset):
+    "A list of inputs `x` and labels `y` with optional `tfms`."
+    def __init__(self, x:ItemList, y:ItemList, tfms:TfmList=None, tfm_y:bool=False, **kwargs):
+        self.x,self.y,self.tfm_y = x,y,tfm_y
+        self.y.x = x
+        self.item=None
+        self.transform(tfms, **kwargs)
+
+    def __len__(self)->int: return len(self.x) if self.item is None else 1
+
+    @contextmanager
+    def set_item(self,item):
+        "For inference, will briefly replace the dataset with one that only contains `item`."
+        self.item = self.x.process_one(item)
+        yield None
+        self.item = None
+
+    def __repr__(self)->str:
+        items = [self[i] for i in range(min(5,len(self.items)))]
+        res = f'{self.__class__.__name__} ({len(self.items)} items)\n'
+        res += f'x: {self.x.__class__.__name__}\n{show_some([i[0] for i in items])}\n'
+        res += f'y: {self.y.__class__.__name__}\n{show_some([i[1] for i in items])}\n'
+        return res + f'Path: {self.path}'
+
+    def predict(self, res):
+        "Delegates predict call on `res` to `self.y`."
+        return self.y.predict(res)
+
+    @property
+    def c(self): return self.y.c
+
+    def new(self, x, y, tfms=None, tfm_y=None, **kwargs)->'LabelList':
+        tfms,tfm_y = ifnone(tfms, self.tfms),ifnone(tfm_y, self.tfm_y)
+        if isinstance(x, ItemList):
+            return self.__class__(x, y, tfms=tfms, tfm_y=tfm_y, **self.tfmargs)
+        else:
+            return self.new(self.x.new(x, **kwargs), self.y.new(y, **kwargs), tfms=tfms, tfm_y=tfm_y).process()
+
+    def __getattr__(self,k:str)->Any:
+        x = super().__getattribute__('x')
+        res = getattr(x, k, None)
+        if res is not None and k not in ['classes', 'c']: return res
+        y = super().__getattribute__('y')
+        res = getattr(y, k, None)
+        if res is not None: return res
+        raise AttributeError(k)
+
+    def __setstate__(self,data:Any): self.__dict__.update(data)
+
+    def __getitem__(self,idxs:Union[int,np.ndarray])->'LabelList':
+        "return a single (x, y) if `idxs` is an integer or a new `LabelList` object if `idxs` is a range."
+        idxs = try_int(idxs)
+        if isinstance(idxs, Integral):
+            if self.item is None: x,y = self.x[idxs],self.y[idxs]
+            else:                 x,y = self.item   ,0
+            if self.tfms or self.tfmargs:
+                x = x.apply_tfms(self.tfms, is_x=True, **self.tfmargs)
+            if hasattr(self, 'tfms_y') and self.tfm_y and self.item is None:
+                y = y.apply_tfms(self.tfms_y, is_x=False, **{**self.tfmargs_y, 'do_resolve':False})
+            if y is None: y=0
+            return x,y
+        else: return self.new(self.x[idxs], self.y[idxs])
+
+    def to_df(self)->None:
+        "Create `pd.DataFrame` containing `items` from `self.x` and `self.y`."
+        return pd.DataFrame(dict(x=self.x._relative_item_paths(), y=[str(o) for o in self.y]))
+
+    def to_csv(self, dest:str)->None:
+        "Save `self.to_df()` to a CSV file in `self.path`/`dest`."
+        self.to_df().to_csv(self.path/dest, index=False)
+
+    def get_state(self, **kwargs):
+        "Return the minimal state for export."
+        state = {'x_cls':self.x.__class__, 'x_proc':self.x.processor,
+                 'y_cls':self.y.__class__, 'y_proc':self.y.processor,
+                 'tfms':self.tfms, 'tfm_y':self.tfm_y, 'tfmargs':self.tfmargs}
+        if hasattr(self, 'tfms_y'):    state['tfms_y']    = self.tfms_y
+        if hasattr(self, 'tfmargs_y'): state['tfmargs_y'] = self.tfmargs_y
+        return {**state, **kwargs}
+
+    def export(self, fn:PathOrStr, **kwargs):
+        "Export the minimal state and save it in `fn` to load an empty version for inference."
+        pickle.dump(self.get_state(**kwargs), open(fn, 'wb'))
+
+    @classmethod
+    def load_empty(cls, path:PathOrStr, fn:PathOrStr):
+        "Load the state in `fn` to create an empty `LabelList` for inference."
+        return cls.load_state(path, pickle.load(open(Path(path)/fn, 'rb')))
+
+    @classmethod
+    def load_state(cls, path:PathOrStr, state:dict) -> 'LabelList':
+        "Create a `LabelList` from `state`."
+        x = state['x_cls']([], path=path, processor=state['x_proc'], ignore_empty=True)
+        y = state['y_cls']([], path=path, processor=state['y_proc'], ignore_empty=True)
+        res = cls(x, y, tfms=state['tfms'], tfm_y=state['tfm_y'], **state['tfmargs']).process()
+        if state.get('tfms_y', False):    res.tfms_y    = state['tfms_y']
+        if state.get('tfmargs_y', False): res.tfmargs_y = state['tfmargs_y']
+        if state.get('normalize', False): res.normalize = state['normalize']
+        return res
+
+    def process(self, xp:PreProcessor=None, yp:PreProcessor=None, name:str=None):
+        "Launch the processing on `self.x` and `self.y` with `xp` and `yp`."
+        self.y.process(yp)
+        if getattr(self.y, 'filter_missing_y', False):
+            filt = array([o is None for o in self.y.items])
+            if filt.sum()>0:
+                #Warnings are given later since progress_bar might make them disappear.
+                self.warn = f"You are labelling your items with {self.y.__class__.__name__}.\n"
+                self.warn += f"Your {name} set contained the following unknown labels, the corresponding items have been discarded.\n"
+                for p in self.y.processor:
+                    if len(getattr(p, 'warns', [])) > 0:
+                        warnings = list(set(p.warns))
+                        self.warn += ', '.join(warnings[:5])
+                        if len(warnings) > 5: self.warn += "..."
+                    p.warns = []
+                self.x,self.y = self.x[~filt],self.y[~filt]
+        self.x.process(xp)
+        return self
+
+    def filter_by_func(self, func:Callable):
+        filt = array([func(x,y) for x,y in zip(self.x.items, self.y.items)])
+        self.x,self.y = self.x[~filt],self.y[~filt]
+        return self
+
+    def transform(self, tfms:TfmList, tfm_y:bool=None, **kwargs):
+        "Set the `tfms` and `tfm_y` value to be applied to the inputs and targets."
+        _check_kwargs(self.x, tfms, **kwargs)
+        if tfm_y is None: tfm_y = self.tfm_y
+        tfms_y = None if tfms is None else list(filter(lambda t: getattr(t, 'use_on_y', True), listify(tfms)))
+        if tfm_y: _check_kwargs(self.y, tfms_y, **kwargs)
+        self.tfms,self.tfmargs  = tfms,kwargs
+        self.tfm_y,self.tfms_y,self.tfmargs_y = tfm_y,tfms_y,kwargs
+        return self
+
+    def transform_y(self, tfms:TfmList=None, **kwargs):
+        "Set `tfms` to be applied to the targets only."
+        tfms_y = list(filter(lambda t: getattr(t, 'use_on_y', True), listify(self.tfms if tfms is None else tfms)))
+        tfmargs_y = {**self.tfmargs, **kwargs} if tfms is None else kwargs
+        _check_kwargs(self.y, tfms_y, **tfmargs_y)
+        self.tfm_y,self.tfms_y,self.tfmargs_y=True,tfms_y,tfmargs_y
+        return self
+
+    def databunch(self, **kwargs):
+        "To throw a clear error message when the data wasn't split."
+        raise Exception("Your data isn't split, if you don't want a validation set, please use `split_none`")
+
+@classmethod
+def _databunch_load_empty(cls, path, fname:str='export.pkl'):
+    "Load an empty `DataBunch` from the exported file in `path/fname` with optional `tfms`."
+    sd = LabelLists.load_empty(path, fn=fname)
+    return sd.databunch()
+
+DataBunch.load_empty = _databunch_load_empty
+
+class MixedProcessor(PreProcessor):
+    def __init__(self, procs:Collection[Union[PreProcessor, Collection[PreProcessor]]]):
+        self.procs = procs
+
+    def process_one(self, item:Any):
+        res = []
+        for procs, i in zip(self.procs, item):
+            for p in procs: i = p.process_one(i)
+            res.append(i)
+        return res
+
+    def process(self, ds:Collection):
+        for procs, il in zip(self.procs, ds.item_lists):
+            for p in procs: p.process(il)
+
+class MixedItem(ItemBase):
+    def __init__(self, items):
+        self.obj = items
+        self.data = [item.data for item in items]
+
+    def __repr__(self): return '\n'.join([f'{self.__class__.__name__}'] + [repr(item) for item in self.obj])
+
+    def apply_tfms(self, tfms:Collection, **kwargs):
+        self.obj = [item.apply_tfms(t, **kwargs) for item,t in zip(self.obj, tfms)]
+        self.data = [item.data for item in self.obj]
+        return self
+
+class MixedItemList(ItemList):
+
+    def __init__(self, item_lists, path:PathOrStr=None, label_cls:Callable=None, inner_df:Any=None,
+                 x:'ItemList'=None, ignore_empty:bool=False, processor=None):
+        self.item_lists = item_lists
+        if processor is None:
+            default_procs = [[p(ds=il) for p in listify(il._processor)] for il in item_lists]
+            processor = MixedProcessor([ifnone(il.processor, dp) for il,dp in zip(item_lists, default_procs)])
+        items = range_of(item_lists[0]) if len(item_lists) >= 1 else []
+        if path is None and len(item_lists) >= 1: path = item_lists[0].path
+        super().__init__(items, processor=processor, path=path,
+                         label_cls=label_cls, inner_df=inner_df, x=x, ignore_empty=ignore_empty)
+
+    def new(self, item_lists, processor:PreProcessor=None, **kwargs)->'ItemList':
+        "Create a new `ItemList` from `items`, keeping the same attributes."
+        processor = ifnone(processor, self.processor)
+        copy_d = {o:getattr(self,o) for o in self.copy_new}
+        kwargs = {**copy_d, **kwargs}
+        return self.__class__(item_lists, processor=processor, **kwargs)
+
+    def get(self, i):
+        return MixedItem([il.get(i) for il in self.item_lists])
+
+    def __getitem__(self,idxs:int)->Any:
+        idxs = try_int(idxs)
+        if isinstance(idxs, Integral): return self.get(idxs)
+        else:
+            item_lists = [il.new(il.items[idxs], inner_df=index_row(il.inner_df, idxs)) for il in self.item_lists]
+            return self.new(item_lists, inner_df=index_row(self.inner_df, idxs))

DeOldify/fastai/datasets.py

@@ -0,0 +1,234 @@
+from .core import *
+import hashlib
+
+__all__ = ['URLs', 'Config', 'untar_data', 'download_data', 'datapath4file', 'url2name', 'url2path']
+
+MODEL_URL = 'http://files.fast.ai/models/'
+URL = 'http://files.fast.ai/data/examples/'
+class URLs():
+    "Global constants for dataset and model URLs."
+    LOCAL_PATH = Path.cwd()
+    S3 = 'https://s3.amazonaws.com/fast-ai-'
+
+    S3_IMAGE    = f'{S3}imageclas/'
+    S3_IMAGELOC = f'{S3}imagelocal/'
+    S3_NLP      = f'{S3}nlp/'
+    S3_COCO     = f'{S3}coco/'
+    S3_MODEL    = f'{S3}modelzoo/'
+
+    # main datasets
+    ADULT_SAMPLE        = f'{URL}adult_sample'
+    BIWI_SAMPLE         = f'{URL}biwi_sample'
+    CIFAR               = f'{URL}cifar10'
+    COCO_SAMPLE         = f'{S3_COCO}coco_sample'
+    COCO_TINY           = f'{URL}coco_tiny'
+    HUMAN_NUMBERS       = f'{URL}human_numbers'
+    IMDB                = f'{S3_NLP}imdb'
+    IMDB_SAMPLE         = f'{URL}imdb_sample'
+    ML_SAMPLE           = f'{URL}movie_lens_sample'
+    MNIST_SAMPLE        = f'{URL}mnist_sample'
+    MNIST_TINY          = f'{URL}mnist_tiny'
+    MNIST_VAR_SIZE_TINY = f'{S3_IMAGE}mnist_var_size_tiny'
+    PLANET_SAMPLE       = f'{URL}planet_sample'
+    PLANET_TINY         = f'{URL}planet_tiny'
+    IMAGENETTE          = f'{S3_IMAGE}imagenette'
+    IMAGENETTE_160      = f'{S3_IMAGE}imagenette-160'
+    IMAGENETTE_320      = f'{S3_IMAGE}imagenette-320'
+    IMAGEWOOF           = f'{S3_IMAGE}imagewoof'
+    IMAGEWOOF_160       = f'{S3_IMAGE}imagewoof-160'
+    IMAGEWOOF_320       = f'{S3_IMAGE}imagewoof-320'
+
+    # kaggle competitions download dogs-vs-cats -p {DOGS.absolute()}
+    DOGS = f'{URL}dogscats'
+
+    # image classification datasets
+    CALTECH_101  = f'{S3_IMAGE}caltech_101'
+    CARS         = f'{S3_IMAGE}stanford-cars'
+    CIFAR_100    = f'{S3_IMAGE}cifar100'
+    CUB_200_2011 = f'{S3_IMAGE}CUB_200_2011'
+    FLOWERS      = f'{S3_IMAGE}oxford-102-flowers'
+    FOOD         = f'{S3_IMAGE}food-101'
+    MNIST        = f'{S3_IMAGE}mnist_png'
+    PETS         = f'{S3_IMAGE}oxford-iiit-pet'
+
+    # NLP datasets
+    AG_NEWS                 = f'{S3_NLP}ag_news_csv'
+    AMAZON_REVIEWS          = f'{S3_NLP}amazon_review_full_csv'
+    AMAZON_REVIEWS_POLARITY = f'{S3_NLP}amazon_review_polarity_csv'
+    DBPEDIA                 = f'{S3_NLP}dbpedia_csv'
+    MT_ENG_FRA              = f'{S3_NLP}giga-fren'
+    SOGOU_NEWS              = f'{S3_NLP}sogou_news_csv'
+    WIKITEXT                = f'{S3_NLP}wikitext-103'
+    WIKITEXT_TINY           = f'{S3_NLP}wikitext-2'
+    YAHOO_ANSWERS           = f'{S3_NLP}yahoo_answers_csv'
+    YELP_REVIEWS            = f'{S3_NLP}yelp_review_full_csv'
+    YELP_REVIEWS_POLARITY   = f'{S3_NLP}yelp_review_polarity_csv'
+
+    # Image localization datasets
+    BIWI_HEAD_POSE     = f"{S3_IMAGELOC}biwi_head_pose"
+    CAMVID             = f'{S3_IMAGELOC}camvid'
+    CAMVID_TINY        = f'{URL}camvid_tiny'
+    LSUN_BEDROOMS      = f'{S3_IMAGE}bedroom'
+    PASCAL_2007        = f'{S3_IMAGELOC}pascal_2007'
+    PASCAL_2012        = f'{S3_IMAGELOC}pascal_2012'
+
+    #Pretrained models
+    OPENAI_TRANSFORMER = f'{S3_MODEL}transformer'
+    WT103_FWD          = f'{S3_MODEL}wt103-fwd'
+    WT103_BWD          = f'{S3_MODEL}wt103-bwd'
+
+# to create/update a checksum for ./mnist_var_size_tiny.tgz, run:
+# python -c 'import fastai.datasets; print(fastai.datasets._check_file("mnist_var_size_tiny.tgz"))'
+_checks = {
+    URLs.ADULT_SAMPLE:(968212, '64eb9d7e23732de0b138f7372d15492f'),
+    URLs.AG_NEWS:(11784419, 'b86f328f4dbd072486591cb7a5644dcd'),
+    URLs.AMAZON_REVIEWS_POLARITY:(688339454, '676f7e5208ec343c8274b4bb085bc938'),
+    URLs.AMAZON_REVIEWS:(643695014, '4a1196cf0adaea22f4bc3f592cddde90'),
+    URLs.BIWI_HEAD_POSE:(452316199, '00f4ccf66e8cba184bc292fdc08fb237'),
+    URLs.BIWI_SAMPLE:(593774, '9179f4c1435f4b291f0d5b072d60c2c9'),
+    URLs.CALTECH_101:(131740031, 'd673425306e98ee4619fcdeef8a0e876'),
+    URLs.CAMVID:(598913237, '648371e4f3a833682afb39b08a3ce2aa'),
+    URLs.CAMVID_TINY:(2314212, '2cf6daf91b7a2083ecfa3e9968e9d915'),
+    URLs.CARS:(1957803273, '9045d6673c9ced0889f41816f6bf2f9f'),
+    URLs.CIFAR:(168168549, 'a5f8c31371b63a406b23368042812d3c'),
+    URLs.CIFAR_100:(169168619, 'e5e65dcb54b9d3913f7b8a9ad6607e62'),
+    URLs.COCO_SAMPLE:(3245877008, '006cd55d633d94b36ecaf661467830ec'),
+    URLs.COCO_TINY:(801038, '367467451ac4fba79a647753c2c66d3a'),
+    URLs.CUB_200_2011:(1150585339, 'd2acaa99439dff0483c7bbac1bfe2a92'),
+    URLs.DBPEDIA:(68341743, '239c7837b9e79db34486f3de6a00e38e'),
+    URLs.DOGS:(839285364, '3e483c8d6ef2175e9d395a6027eb92b7'),
+    URLs.FLOWERS:(345236087, '5666e01c1311b4c67fcf20d2b3850a88'),
+    URLs.FOOD:(5686607260, '1a540ebf1fb40b2bf3f2294234ba7907'),
+    URLs.HUMAN_NUMBERS:(30252, '8a19c3bfa2bcb08cd787e741261f3ea2'),
+    URLs.IMDB:(144440600, '90f9b1c4ff43a90d67553c9240dc0249'),
+    URLs.IMDB_SAMPLE:(571827, '0842e61a9867caa2e6fbdb14fa703d61'),
+    URLs.LSUN_BEDROOMS:(4579163978, '35d84f38f8a15fe47e66e460c8800d68'),
+    URLs.ML_SAMPLE:(51790, '10961384dfe7c5181460390a460c1f77'),
+    URLs.MNIST:(15683414, '03639f83c4e3d19e0a3a53a8a997c487'),
+    URLs.MNIST_SAMPLE:(3214948, '2dbc7ec6f9259b583af0072c55816a88'),
+    URLs.MNIST_TINY:(342207, '56143e8f24db90d925d82a5a74141875'),
+    URLs.MNIST_VAR_SIZE_TINY:(565372, 'b71a930f4eb744a4a143a6c7ff7ed67f'),
+    URLs.MT_ENG_FRA:(2598183296, '69573f58e2c850b90f2f954077041d8c'),
+    URLs.OPENAI_TRANSFORMER:(432848315, '024b0d2203ebb0cd1fc64b27cf8af18e'),
+    URLs.PASCAL_2007:(1636130334, 'a70574e9bc592bd3b253f5bf46ce12e3'),
+    URLs.PASCAL_2012:(2611715776, '2ae7897038383836f86ce58f66b09e31'),
+    URLs.PETS:(811706944, 'e4db5c768afd933bb91f5f594d7417a4'),
+    URLs.PLANET_SAMPLE:(15523994, '8bfb174b3162f07fbde09b54555bdb00'),
+    URLs.PLANET_TINY:(997569, '490873c5683454d4b2611fb1f00a68a9'),
+    URLs.SOGOU_NEWS:(384269937, '950f1366d33be52f5b944f8a8b680902'),
+    URLs.WIKITEXT:(190200704, '2dd8cf8693b3d27e9c8f0a7df054b2c7'),
+    URLs.WIKITEXT_TINY:(4070055, '2a82d47a7b85c8b6a8e068dc4c1d37e7'),
+    URLs.WT103_FWD:(105067061, '7d1114cd9684bf9d1ca3c9f6a54da6f9'),
+    URLs.WT103_BWD:(105205312, '20b06f5830fd5a891d21044c28d3097f'),
+    URLs.YAHOO_ANSWERS:(319476345, '0632a0d236ef3a529c0fa4429b339f68'),
+    URLs.YELP_REVIEWS_POLARITY:(166373201, '48c8451c1ad30472334d856b5d294807'),
+    URLs.YELP_REVIEWS:(196146755, '1efd84215ea3e30d90e4c33764b889db'),
+}
+
+#TODO: This can probably be coded more shortly and nicely.
+class Config():
+    "Creates a default config file 'config.yml' in $FASTAI_HOME (default `~/.fastai/`)"
+    DEFAULT_CONFIG_LOCATION = os.path.expanduser(os.getenv('FASTAI_HOME', '~/.fastai'))
+    DEFAULT_CONFIG_PATH = DEFAULT_CONFIG_LOCATION + '/config.yml'
+    DEFAULT_CONFIG = {
+        'data_path': DEFAULT_CONFIG_LOCATION + '/data',
+        'data_archive_path': DEFAULT_CONFIG_LOCATION + '/data',
+        'model_path': DEFAULT_CONFIG_LOCATION + '/models'
+    }
+
+    @classmethod
+    def get_key(cls, key):
+        "Get the path to `key` in the config file."
+        return cls.get().get(key, cls.DEFAULT_CONFIG.get(key,None))
+
+    @classmethod
+    def get_path(cls, path):
+        "Get the `path` in the config file."
+        return _expand_path(cls.get_key(path))
+
+    @classmethod
+    def data_path(cls):
+        "Get the path to data in the config file."
+        return cls.get_path('data_path')
+
+    @classmethod
+    def data_archive_path(cls):
+        "Get the path to data archives in the config file."
+        return cls.get_path('data_archive_path')
+
+    @classmethod
+    def model_path(cls):
+        "Get the path to fastai pretrained models in the config file."
+        return cls.get_path('model_path')
+
+    @classmethod
+    def get(cls, fpath=None, create_missing=True):
+        "Retrieve the `Config` in `fpath`."
+        fpath = _expand_path(fpath or cls.DEFAULT_CONFIG_PATH)
+        if not fpath.exists() and create_missing: cls.create(fpath)
+        assert fpath.exists(), f'Could not find config at: {fpath}. Please create'
+        with open(fpath, 'r') as yaml_file: return yaml.safe_load(yaml_file)
+
+    @classmethod
+    def create(cls, fpath):
+        "Creates a `Config` from `fpath`."
+        fpath = _expand_path(fpath)
+        assert(fpath.suffix == '.yml')
+        if fpath.exists(): return
+        fpath.parent.mkdir(parents=True, exist_ok=True)
+        with open(fpath, 'w') as yaml_file:
+            yaml.dump(cls.DEFAULT_CONFIG, yaml_file, default_flow_style=False)
+
+def _expand_path(fpath): return Path(fpath).expanduser()
+def url2name(url): return url.split('/')[-1]
+
+#TODO: simplify this mess
+def url2path(url, data=True, ext:str='.tgz'):
+    "Change `url` to a path."
+    name = url2name(url)
+    return datapath4file(name, ext=ext, archive=False) if data else modelpath4file(name, ext=ext)
+def _url2tgz(url, data=True, ext:str='.tgz'):
+    return datapath4file(f'{url2name(url)}{ext}', ext=ext) if data else modelpath4file(f'{url2name(url)}{ext}', ext=ext)
+
+def modelpath4file(filename, ext:str='.tgz'):
+    "Return model path to `filename`, checking locally first then in the config file."
+    local_path = URLs.LOCAL_PATH/'models'/filename
+    if local_path.exists() or local_path.with_suffix(ext).exists(): return local_path
+    else: return Config.model_path()/filename
+
+def datapath4file(filename, ext:str='.tgz', archive=True):
+    "Return data path to `filename`, checking locally first then in the config file."
+    local_path = URLs.LOCAL_PATH/'data'/filename
+    if local_path.exists() or local_path.with_suffix(ext).exists(): return local_path
+    elif archive: return Config.data_archive_path() / filename
+    else: return Config.data_path() / filename
+
+def download_data(url:str, fname:PathOrStr=None, data:bool=True, ext:str='.tgz') -> Path:
+    "Download `url` to destination `fname`."
+    fname = Path(ifnone(fname, _url2tgz(url, data, ext=ext)))
+    os.makedirs(fname.parent, exist_ok=True)
+    if not fname.exists():
+        print(f'Downloading {url}')
+        download_url(f'{url}{ext}', fname)
+    return fname
+
+def _check_file(fname):
+    size = os.path.getsize(fname)
+    with open(fname, "rb") as f:
+        hash_nb = hashlib.md5(f.read(2**20)).hexdigest()
+    return size,hash_nb
+
+def untar_data(url:str, fname:PathOrStr=None, dest:PathOrStr=None, data=True, force_download=False) -> Path:
+    "Download `url` to `fname` if `dest` doesn't exist, and un-tgz to folder `dest`."
+    dest = url2path(url, data) if dest is None else Path(dest)/url2name(url)
+    fname = Path(ifnone(fname, _url2tgz(url, data)))
+    if force_download or (fname.exists() and url in _checks and _check_file(fname) != _checks[url]):
+        print(f"A new version of the {'dataset' if data else 'model'} is available.")
+        if fname.exists(): os.remove(fname)
+        if dest.exists(): shutil.rmtree(dest)
+    if not dest.exists():
+        fname = download_data(url, fname=fname, data=data)
+        if url in _checks:
+            assert _check_file(fname) == _checks[url], f"Downloaded file {fname} does not match checksum expected! Remove that file from {Config().data_archive_path()} and try your code again."
+        tarfile.open(fname, 'r:gz').extractall(dest.parent)
+    return dest

DeOldify/fastai/distributed.py

@@ -0,0 +1,119 @@
+from .torch_core import *
+from .basic_train import Learner,LearnerCallback
+from torch.nn.parallel import DistributedDataParallel, DataParallel
+from torch.utils.data.distributed import DistributedSampler
+
+from fastai.text import TextLMDataBunch
+
+__all__ = ['DistributedRecorder', 'DistributedTrainer', 'read_metrics', 'setup_distrib']
+
+def rnn_reset(self):
+    if hasattr(self.module, 'reset'): self.module.reset()
+DistributedDataParallel.reset = rnn_reset
+
+class ParallelTrainer(LearnerCallback):
+    _order = -20
+    def on_train_begin(self, **kwargs): self.learn.model = DataParallel(self.learn.model)
+    def on_train_end  (self, **kwargs): self.learn.model = self.learn.model.module
+
+class DistributedTrainer(LearnerCallback):
+    _order = -20 # Needs to run before the recorder
+    def __init__(self, learn:Learner, cuda_id:int=0):
+        super().__init__(learn)
+        self.cuda_id,self.train_sampler = cuda_id,None
+
+    def _change_dl(self, dl, shuffle):
+        old_dl = dl
+        sampler = OurDistributedSampler(dl.dataset, shuffle=shuffle)
+        new_dl = dl.new(shuffle=False, sampler=sampler)
+        return old_dl,new_dl,sampler
+
+    def on_train_begin(self, **kwargs):
+        self.learn.model = DistributedDataParallel(self.model, device_ids=[self.cuda_id], output_device=self.cuda_id)
+        shuffle = self.data.train_dl.init_kwargs['shuffle'] if hasattr(self.data.train_dl, 'init_kwargs') else True
+        self.old_train_dl,self.data.train_dl,self.train_sampler = self._change_dl(self.data.train_dl, shuffle)
+        if hasattr(self.data, 'valid_dl') and self.data.valid_dl is not None:
+            self.old_valid_dl,self.data.valid_dl,self.valid_sampler = self._change_dl(self.data.valid_dl, shuffle)
+        self.rank = rank_distrib()
+        self.recorder.silent = (self.rank != 0)
+
+    def on_epoch_begin(self, epoch, **kwargs): self.train_sampler.set_epoch(epoch)
+
+    def on_train_end(self, **kwargs):
+        self.learn.model = self.learn.model.module
+        self.learn.data.train_dl = self.old_train_dl
+        if hasattr(self.learn.data, 'valid_dl') and self.learn.data.valid_dl is not None:
+            self.learn.data.valid_dl = self.old_valid_dl
+
+class DistributedRecorder(LearnerCallback):
+    def __init__(self, learn:Learner, cuda_id:int=0, cache_dir:PathOrStr='tmp'):
+        super().__init__(learn)
+        self.cuda_id,self.cache_dir = cuda_id,cache_dir
+
+    def on_train_begin(self, **kwargs):
+        os.makedirs(self.learn.path/self.cache_dir, exist_ok=True)
+
+    def on_epoch_end(self, **kwargs): self.save_stats()
+    def on_train_end(self, **kwargs): self.save_stats()
+
+    def save_stats(self):
+        cache_path,recorder = self.learn.path/self.cache_dir,self.learn.recorder
+        np.save(cache_path/f'losses_{self.cuda_id}', np.array(recorder.losses))
+        stats = np.array([[v] + m for v,m in zip(recorder.val_losses,recorder.metrics)])
+        np.save(cache_path/f'metrics_{self.cuda_id}', stats)
+
+def _learner_parallel(learn:Learner):
+    "Use nn.DataParallel when training and remove when done"
+    if not torch.cuda.is_available(): warnings.warn('CUDA is not available, check your drivers - training will continue on CPU', ResourceWarning) 
+    learn.callbacks.append(ParallelTrainer(learn))
+    return learn
+
+def _learner_distributed(learn:Learner, cuda_id:int, cache_dir:PathOrStr='tmp'):
+    "Put `learn` on distributed training with `cuda_id`."
+    learn.callbacks.append(DistributedTrainer(learn, cuda_id))
+    learn.callbacks.append(DistributedRecorder(learn, cuda_id, cache_dir))
+    return learn
+
+Learner.to_distributed = _learner_distributed
+Learner.to_parallel = _learner_parallel
+
+def read_metrics(cache_path:PathOrStr, n_gpus:int, reduce:bool=True):
+    losses,metrics = [],[]
+    for i in range(n_gpus):
+        losses.append(np.load(cache_path/f'losses_{i}.npy')[None])
+        metrics.append(np.load(cache_path/f'metrics_{i}.npy')[None])
+    if reduce:
+        losses,metrics = np.concatenate(losses,0),np.concatenate(metrics,0)
+        return losses.mean(0),metrics.mean(0)
+    return losses,metrics
+
+def setup_distrib(gpu:Any=None):
+    if gpu is None: return gpu
+    gpu = int(gpu)
+    torch.cuda.set_device(int(gpu))
+    if num_distrib() > 1:
+        torch.distributed.init_process_group(backend='nccl', init_method='env://')
+    return gpu
+
+class OurDistributedSampler(DistributedSampler):
+    "A sampler for language models with the option to not shuffle."
+    def __init__(self, dataset, num_replicas=None, rank=None, shuffle=True):
+            super().__init__(dataset, num_replicas=num_replicas, rank=rank)
+            self.shuffle = shuffle
+    
+    def __iter__(self):
+        if self.shuffle:
+            g = torch.Generator()
+            g.manual_seed(self.epoch)
+            indices = torch.randperm(len(self.dataset), generator=g).tolist()
+        else: indices = torch.arange(len(self.dataset)).tolist()
+
+        # add extra samples to make it evenly divisible
+        indices += indices[:(self.total_size - len(indices))]
+        assert len(indices) == self.total_size
+
+        # subsample
+        indices = indices[self.rank:self.total_size:self.num_replicas]
+        assert len(indices) == self.num_samples
+
+        return iter(indices)

DeOldify/fastai/gen_doc/__init__.py

@@ -0,0 +1 @@
+from . import gen_notebooks, nbdoc, core, doctest, nbtest

DeOldify/fastai/gen_doc/autogen.tpl

@@ -0,0 +1,10 @@
+<!--
+
+
+#################################################
+### THIS FILE WAS AUTOGENERATED! DO NOT EDIT! ###
+#################################################
+# file to edit: {{ resources.nb_path }}
+# instructions: https://docs.fast.ai/gen_doc_main.html
+
+-->

DeOldify/fastai/gen_doc/convert2html.py

@@ -0,0 +1,51 @@
+import os.path, re, nbformat, jupyter_contrib_nbextensions
+from nbconvert.preprocessors import Preprocessor
+from nbconvert import HTMLExporter
+from traitlets.config import Config
+from pathlib import Path
+
+__all__ = ['read_nb', 'convert_nb', 'convert_all']
+
+exporter = HTMLExporter(Config())
+exporter.exclude_input_prompt=True
+exporter.exclude_output_prompt=True
+#Loads the template to deal with hidden cells.
+exporter.template_file = 'jekyll.tpl'
+path = Path(__file__).parent
+exporter.template_path.append(str(path))
+
+def read_nb(fname):
+    "Read the notebook in `fname`."
+    with open(fname,'r') as f: return nbformat.reads(f.read(), as_version=4)
+
+def convert_nb(fname, dest_path='.'):
+    "Convert a notebook `fname` to html file in `dest_path`."
+    from .gen_notebooks import remove_undoc_cells, remove_code_cell_jupyter_widget_state_elem
+    nb = read_nb(fname)
+    nb['cells'] = remove_undoc_cells(nb['cells'])
+    nb['cells'] = remove_code_cell_jupyter_widget_state_elem(nb['cells'])
+    fname = Path(fname).absolute()
+    dest_name = fname.with_suffix('.html').name
+    meta = nb['metadata']
+    meta_jekyll = meta['jekyll'] if 'jekyll' in meta else {'title': fname.with_suffix('').name}
+    meta_jekyll['nb_path'] = f'{fname.parent.name}/{fname.name}'
+    with open(f'{dest_path}/{dest_name}','w') as f:
+        f.write(exporter.from_notebook_node(nb, resources=meta_jekyll)[0])
+
+def convert_all(folder, dest_path='.', force_all=False):
+    "Convert modified notebooks in `folder` to html pages in `dest_path`."
+    path = Path(folder)
+
+    changed_cnt = 0
+    for fname in path.glob("*.ipynb"):
+        # only rebuild modified files
+        fname_out = Path(dest_path)/fname.with_suffix('.html').name
+        if not force_all and fname_out.exists():
+            in_mod  = os.path.getmtime(fname)
+            out_mod = os.path.getmtime(fname_out)
+            if in_mod < out_mod: continue
+
+        print(f"converting: {fname} => {fname_out}")
+        changed_cnt += 1
+        convert_nb(fname, dest_path=dest_path)
+    if not changed_cnt: print("No notebooks were modified")

DeOldify/fastai/gen_doc/core.py

@@ -0,0 +1,5 @@
+from ..core import *
+import re
+
+def strip_fastai(s):  return re.sub(r'^fastai\.', '', s)
+

DeOldify/fastai/gen_doc/docstrings.py

@@ -0,0 +1,142 @@
+# https://github.com/openstack/rally/blob/master/rally/common/plugin/info.py
+# Copyright 2015: Mirantis Inc.
+# All Rights Reserved.
+#
+#    Licensed under the Apache License, Version 2.0 (the "License"); you may
+#    not use this file except in compliance with the License. You may obtain
+#    a copy of the License at
+#
+#         http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+#    License for the specific language governing permissions and limitations
+#    under the License.
+
+import re
+import sys
+
+__all__ = ['parse_docstring']
+
+
+FIELDS = 'param|val' # supported fields
+PARAM_OR_RETURN_REGEX = re.compile(f":(?:{FIELDS}|return)")
+RETURN_REGEX = re.compile(":return: (?P<doc>.*)", re.S)
+NEW_REGEX = re.compile(f":(?P<field>{FIELDS}) (?P<name>[\*\w]+): (?P<doc>.*?)"
+                         f"(?:(?=:(?:{FIELDS}|return|raises))|\Z)", re.S)
+
+def trim(docstring):
+    """trim function from PEP-257"""
+    if not docstring:
+        return ""
+    # Convert tabs to spaces (following the normal Python rules)
+    # and split into a list of lines:
+    lines = docstring.expandtabs().splitlines()
+    # Determine minimum indentation (first line doesn't count):
+    indent = sys.maxsize
+    for line in lines[1:]:
+        stripped = line.lstrip()
+        if stripped:
+            indent = min(indent, len(line) - len(stripped))
+    # Remove indentation (first line is special):
+    trimmed = [lines[0].strip()]
+    if indent < sys.maxsize:
+        for line in lines[1:]:
+            trimmed.append(line[indent:].rstrip())
+    # Strip off trailing and leading blank lines:
+    while trimmed and not trimmed[-1]:
+        trimmed.pop()
+    while trimmed and not trimmed[0]:
+        trimmed.pop(0)
+
+    # Current code/unittests expects a line return at
+    # end of multiline docstrings
+    # workaround expected behavior from unittests
+    if "\n" in docstring:
+        trimmed.append("")
+
+    # Return a single string:
+    return "\n".join(trimmed)
+
+
+def reindent(string):
+    return "\n".join(l.strip() for l in string.strip().split("\n"))
+
+
+def parse_docstring(docstring):
+    """Parse the docstring into its components.
+
+    :return: a dictionary of form
+              {
+                  "short_description": ...,
+                  "long_description": ...,
+                  "params": [{"name": ..., "doc": ...}, ...],
+                  "vals": [{"name": ..., "doc": ...}, ...],
+                  "return": ...
+              }
+    """
+
+    short_description = long_description = return_str = ""
+    args = []
+
+    if docstring:
+        docstring = trim(docstring.lstrip("\n"))
+
+        lines = docstring.split("\n", 1)
+        short_description = lines[0]
+
+        if len(lines) > 1:
+            long_description = lines[1].strip()
+
+            params_return_desc = None
+
+            match = PARAM_OR_RETURN_REGEX.search(long_description)
+            if match:
+                long_desc_end = match.start()
+                params_return_desc = long_description[long_desc_end:].strip()
+                long_description = long_description[:long_desc_end].rstrip()
+
+            if params_return_desc:
+                args = [
+                    {"name": name, "doc": trim(doc), "field": field}
+                    for field, name, doc in NEW_REGEX.findall(params_return_desc)
+                ]
+                match = RETURN_REGEX.search(params_return_desc)
+                if match:
+                    return_str = reindent(match.group("doc"))
+    comments = {p['name']: p['doc'] for p in args}
+    return {
+        "short_description": short_description,
+        "long_description": long_description,
+        "args": args,
+        "comments": comments,
+        "return": return_str
+    }
+
+
+class InfoMixin(object):
+
+    @classmethod
+    def _get_doc(cls):
+        """Return documentary of class
+
+        By default it returns docstring of class, but it can be overridden
+        for example for cases like merging own docstring with parent
+        """
+        return cls.__doc__
+
+    @classmethod
+    def get_info(cls):
+        doc = parse_docstring(cls._get_doc())
+
+        return {
+            "name": cls.get_name(),
+            "platform": cls.get_platform(),
+            "module": cls.__module__,
+            "title": doc["short_description"],
+            "description": doc["long_description"],
+            "parameters": doc["params"],
+            "schema": getattr(cls, "CONFIG_SCHEMA", None),
+            "return": doc["return"]
+        }

DeOldify/fastai/gen_doc/doctest.py

@@ -0,0 +1,133 @@
+import sys, re, json, pprint
+from pathlib import Path
+from collections import defaultdict
+from inspect import currentframe, getframeinfo, ismodule
+
+__all__ = ['this_tests']
+
+DB_NAME = 'test_registry.json'
+
+def _json_set_default(obj):
+    if isinstance(obj, set): return list(obj)
+    raise TypeError
+
+class TestRegistry:
+    "Tests register which API they validate using this class."
+    registry = defaultdict(list)
+    this_tests_check = None
+    missing_this_tests = set()
+
+    # logic for checking whether each test calls `this_tests`:
+    # 1. `this_tests_check` is set to True during test's 'setup' stage if it wasn't skipped
+    # 2. if the test is dynamically skipped `this_tests_check` is set to False
+    # 3. `this_tests` sets this flag to False when it's successfully completes
+    # 4. if during the 'teardown' stage `this_tests_check` is still True then we
+    # know that this test needs `this_tests_check`
+
+    @staticmethod
+    def this_tests(*funcs):
+        prev_frame = currentframe().f_back.f_back
+        file_name, lineno, test_name, _, _ = getframeinfo(prev_frame)
+        parent_func_lineno, _ = get_parent_func(lineno, get_lines(file_name))
+        entry = {'file': relative_test_path(file_name), 'test': test_name , 'line': parent_func_lineno}
+        for func in funcs:
+            if func == 'na':
+                # special case when we can't find a function to declare, e.g.
+                # when attributes are tested
+                continue
+            try:
+                func_fq = get_func_fq_name(func)
+            except:
+                raise Exception(f"'{func}' is not a function") from None
+            if re.match(r'fastai\.', func_fq):
+                if entry not in TestRegistry.registry[func_fq]:
+                    TestRegistry.registry[func_fq].append(entry)
+            else:
+                raise Exception(f"'{func}' is not in the fastai API") from None
+        TestRegistry.this_tests_check = False
+
+    def this_tests_check_on():
+        TestRegistry.this_tests_check = True
+
+    def this_tests_check_off():
+        TestRegistry.this_tests_check = False
+
+    def this_tests_check_run(file_name, test_name):
+        if TestRegistry.this_tests_check:
+            TestRegistry.missing_this_tests.add(f"{file_name}::{test_name}")
+
+    def registry_save():
+        if TestRegistry.registry:
+            path = Path(__file__).parent.parent.resolve()/DB_NAME
+            if path.exists():
+                #print("\n*** Merging with the existing test registry")
+                with open(path, 'r') as f: old_registry = json.load(f)
+                TestRegistry.registry = merge_registries(old_registry, TestRegistry.registry)
+            #print(f"\n*** Saving test registry @ {path}")
+            with open(path, 'w') as f:
+                json.dump(obj=TestRegistry.registry, fp=f, indent=4, sort_keys=True, default=_json_set_default)
+
+    def missing_this_tests_alert():
+        if TestRegistry.missing_this_tests:
+            tests = '\n  '.join(sorted(TestRegistry.missing_this_tests))
+            print(f"""
+*** Attention ***
+Please include `this_tests` call in each of the following tests:
+  {tests}
+For details see: https://docs.fast.ai/dev/test.html#test-registry""")
+
+# merge_registries helpers
+# merge dict of lists of dict
+def a2k(a): return '::'.join([a['file'], a['test']]), a['line']
+def k2a(k, v): f,t = k.split('::'); return {"file": f, "line": v, "test": t}
+# merge by key that is a combination of 2 values: test, file
+def merge_lists(a, b):
+    x = dict(map(a2k, [*a, *b]))            # pack + merge
+    return [k2a(k, v) for k,v in x.items()] # unpack
+def merge_registries(a, b):
+    for i in b: a[i] = merge_lists(a[i], b[i]) if i in a else b[i]
+    return a
+
+def this_tests(*funcs): TestRegistry.this_tests(*funcs)
+
+def str2func(name):
+    "Converts 'fastai.foo.bar' into an function 'object' if such exists"
+    if isinstance(name, str): subpaths = name.split('.')
+    else:                     return None
+
+    module = subpaths.pop(0)
+    if module in sys.modules: obj = sys.modules[module]
+    else:                     return None
+
+    for subpath in subpaths:
+        obj = getattr(obj, subpath, None)
+        if obj == None: return None
+    return obj
+
+def get_func_fq_name(func):
+    if ismodule(func): return func.__name__
+    if isinstance(func, str): func = str2func(func)
+    name = None
+    if   hasattr(func, '__qualname__'): name = func.__qualname__
+    elif hasattr(func, '__name__'):     name = func.__name__
+    elif hasattr(func, '__wrapped__'):  return get_func_fq_name(func.__wrapped__)
+    elif hasattr(func, '__class__'):    name = func.__class__.__name__
+    else: raise Exception(f"'{func}' is not a func or class")
+    return f'{func.__module__}.{name}'
+
+def get_parent_func(lineno, lines, ignore_missing=False):
+    "Find any lines where `elt` is called and return the parent test function"
+    for idx,l in enumerate(reversed(lines[:lineno])):
+        if re.match(f'\s*def test', l):  return (lineno - idx), l # 1 based index for github
+        if re.match(f'\w+', l):  break # top level indent - out of function scope
+    if ignore_missing: return None
+    raise LookupError('Could not find parent function for line:', lineno, lines[:lineno])
+
+def relative_test_path(test_file:Path)->str:
+    "Path relative to the `fastai` parent directory"
+    test_file = Path(test_file)
+    testdir_idx = list(reversed(test_file.parts)).index('tests')
+    return '/'.join(test_file.parts[-(testdir_idx+1):])
+
+def get_lines(file):
+    with open(file, 'r') as f: return f.readlines()

DeOldify/fastai/gen_doc/gen_notebooks.py

@@ -0,0 +1,350 @@
+"`gen_doc.nbdoc` generates notebook documentation from module functions and links to correct places"
+import pkgutil, inspect, sys,os, importlib,json,enum,warnings,nbformat,re
+from IPython.core.display import display, Markdown
+from nbconvert.preprocessors import ExecutePreprocessor
+from nbformat.sign import NotebookNotary
+from pathlib import Path
+from .core import *
+from .nbdoc import *
+
+__all__ = ['create_module_page', 'update_module_page', 'import_mod',
+           'link_nb', 'update_notebooks', 'generate_missing_metadata', 'update_nb_metadata']
+
+def get_empty_notebook():
+    "Default notbook with the minimum metadata."
+    #TODO: check python version and nbformat
+    return {'metadata': {'kernelspec': {'display_name': 'Python 3',
+                                        'language': 'python',
+                                        'name': 'python3'},
+                         'language_info': {'codemirror_mode': {'name': 'ipython', 'version': 3},
+                         'file_extension': '.py',
+                         'mimetype': 'text/x-python',
+                         'name': 'python',
+                         'nbconvert_exporter': 'python',
+                         'pygments_lexer': 'ipython3',
+                         'version': '3.6.6'}},
+            'nbformat': 4,
+            'nbformat_minor': 2}
+
+def get_md_cell(source, metadata=None):
+    "Markdown cell containing `source` with `metadata`."
+    return {'cell_type': 'markdown',
+            'metadata': {} if metadata is None else metadata,
+            'source': source}
+
+def get_empty_cell(ctype='markdown'):
+    "Empty cell of type `ctype`."
+    return {'cell_type': ctype, 'metadata': {}, 'source': []}
+
+def get_code_cell(code, hidden=False):
+    "Code cell containing `code` that may be `hidden`."
+    return {'cell_type' : 'code',
+            'execution_count': 0,
+            'metadata' : {'hide_input': hidden, 'trusted':True},
+            'source' : code,
+            'outputs': []}
+
+def get_doc_cell(func_name):
+    "Code cell with the command to show the doc of `func_name`."
+    code = f"show_doc({func_name})"
+    return get_code_cell(code, True)
+
+def get_global_vars(mod):
+    "Return globally assigned variables."
+    # https://stackoverflow.com/questions/8820276/docstring-for-variable/31764368#31764368
+    import ast,re
+    with open(mod.__file__, 'r') as f: fstr = f.read()
+    flines = fstr.splitlines()
+    d = {}
+    for node in ast.walk(ast.parse(fstr)):
+        if isinstance(node,ast.Assign) and hasattr(node.targets[0], 'id'):
+            key,lineno = node.targets[0].id,node.targets[0].lineno
+            codestr = flines[lineno]
+            match = re.match(f"^({key})\s*=\s*.*", codestr)
+            if match and match.group(1) != '__all__': # only top level assignment
+                d[key] = f'`{codestr}` {get_source_link(mod, lineno)}'
+    return d
+
+def write_nb(nb, nb_path, mode='w'):
+    with open(nb_path, mode) as f: f.write(nbformat.writes(nbformat.from_dict(nb), version=4))
+
+class ExecuteShowDocPreprocessor(ExecutePreprocessor):
+    "An ExecutePreprocessor that only executes show_doc cells"
+    def preprocess_cell(self, cell, resources, index):
+        if 'source' in cell and cell.cell_type == "code":
+            if IMPORT_RE.search(cell['source']) or SHOW_DOC_RE.search(cell['source']):
+                return super().preprocess_cell(cell, resources, index)
+        return cell, resources
+
+def execute_nb(fname, metadata=None, save=True, show_doc_only=False):
+    "Execute notebook `fname` with `metadata` for preprocessing."
+    # Any module used in the notebook that isn't inside must be in the same directory as this script
+    with open(fname) as f: nb = nbformat.read(f, as_version=4)
+    ep_class = ExecuteShowDocPreprocessor if show_doc_only else ExecutePreprocessor
+    ep = ep_class(timeout=600, kernel_name='python3')
+    metadata = metadata or {}
+    ep.preprocess(nb, metadata)
+    if save:
+        with open(fname, 'wt') as f: nbformat.write(nb, f)
+        NotebookNotary().sign(nb)
+
+def _symbol_skeleton(name): return [get_doc_cell(name), get_md_cell(f"`{name}`")]
+
+def create_module_page(mod, dest_path, force=False):
+    "Create the documentation notebook for module `mod_name` in path `dest_path`"
+    nb = get_empty_notebook()
+    mod_name = mod.__name__
+    strip_name = strip_fastai(mod_name)
+    init_cell = [get_md_cell(f'## Title for {strip_name} (use plain english, not module name!)'), get_md_cell('Type an introduction of the package here.')]
+    cells = [get_code_cell(f'from fastai.gen_doc.nbdoc import *\nfrom {mod_name} import * ', True)]
+
+    gvar_map = get_global_vars(mod)
+    if gvar_map: cells.append(get_md_cell('### Global Variable Definitions:'))
+    for name in get_exports(mod):
+        if name in gvar_map: cells.append(get_md_cell(gvar_map[name]))
+
+    for ft_name in get_ft_names(mod, include_inner=True):
+        if not hasattr(mod, ft_name):
+            warnings.warn(f"Module {strip_name} doesn't have a function named {ft_name}.")
+            continue
+        cells += _symbol_skeleton(ft_name)
+        elt = getattr(mod, ft_name)
+    nb['cells'] = init_cell + cells + [get_md_cell(UNDOC_HEADER)]
+
+    doc_path = get_doc_path(mod, dest_path)
+    write_nb(nb, doc_path, 'w' if force else 'x')
+    execute_nb(doc_path)
+    return doc_path
+
+_default_exclude = ['.ipynb_checkpoints', '__pycache__', '__init__.py', 'imports']
+
+def get_module_names(path_dir, exclude=None):
+    if exclude is None: exclude = _default_exclude
+    "Search a given `path_dir` and return all the modules contained inside except those in `exclude`"
+    files = sorted(path_dir.glob('*'), key=lambda x: (x.is_dir(), x.name), reverse=True) # directories first
+    res = [f'{path_dir.name}']
+    for f in files:
+        if f.is_dir() and f.name in exclude: continue # exclude directories
+        if any([f.name.endswith(ex) for ex in exclude]): continue # exclude extensions
+
+        if f.suffix == '.py': res.append(f'{path_dir.name}.{f.stem}')
+        elif f.is_dir(): res += [f'{path_dir.name}.{name}' for name in get_module_names(f)]
+    return res
+
+def read_nb(fname):
+    "Read a notebook in `fname` and return its corresponding json"
+    with open(fname,'r') as f: return nbformat.reads(f.read(), as_version=4)
+
+SHOW_DOC_RE = re.compile(r"show_doc\(([\w\.]*)")
+def read_nb_content(cells, mod_name):
+    "Build a dictionary containing the position of the `cells`."
+    doc_fns = {}
+    for i, cell in enumerate(cells):
+        if cell['cell_type'] == 'code':
+            for match in SHOW_DOC_RE.findall(cell['source']):
+                doc_fns[match] = i
+    return doc_fns
+
+def read_nb_types(cells):
+    doc_fns = {}
+    for i, cell in enumerate(cells):
+        if cell['cell_type'] == 'markdown':
+            match = re.match(r"^(?:<code>|`)?(\w*)\s*=\s*", cell['source'])
+            if match is not None: doc_fns[match.group(1)] = i
+    return doc_fns
+
+def link_markdown_cells(cells, modules):
+    "Create documentation links for all cells in markdown with backticks."
+    for i, cell in enumerate(cells):
+        if cell['cell_type'] == 'markdown':
+            cell['source'] = link_docstring(modules, cell['source'])
+
+def get_insert_idx(pos_dict, name):
+    "Return the position to insert a given function doc in a notebook."
+    keys,i = list(pos_dict.keys()),0
+    while i < len(keys) and str.lower(keys[i]) < str.lower(name): i+=1
+    if i == len(keys): return -1
+    else:              return pos_dict[keys[i]]
+
+def update_pos(pos_dict, start_key, nbr=2):
+    "Update the `pos_dict` by moving all positions after `start_key` by `nbr`."
+    for key,idx in pos_dict.items():
+        if str.lower(key) >= str.lower(start_key): pos_dict[key] += nbr
+    return pos_dict
+
+def insert_cells(cells, pos_dict, ft_name, append=False):
+    "Insert the function doc `cells` at their correct position and updates `pos_dict`."
+    idx = get_insert_idx(pos_dict, ft_name)
+    if append or idx == -1: cells += [get_doc_cell(ft_name), get_empty_cell()]
+    else:
+        cells.insert(idx, get_doc_cell(ft_name))
+        cells.insert(idx+1, get_empty_cell())
+        pos_dict = update_pos(pos_dict, ft_name, 2)
+    return cells, pos_dict
+
+def get_doc_path(mod, dest_path):
+    strip_name = strip_fastai(mod.__name__)
+    return os.path.join(dest_path,f'{strip_name}.ipynb')
+
+def generate_missing_metadata(dest_file):
+    fn = Path(dest_file)
+    meta_fn = fn.parent/'jekyll_metadata.ipynb'
+    if not fn.exists() or not meta_fn.exists(): return print('Could not find notebooks:', fn, meta_fn)
+    metadata_nb = read_nb(meta_fn)
+
+    if has_metadata_cell(metadata_nb['cells'], fn.name): return
+    nb = read_nb(fn)
+    jmd = nb['metadata'].get('jekyll', {})
+    fmt_params = ''
+    for k,v in jmd.items(): fmt_params += f',\n    {k}={stringify(v)}'
+    metadata_cell = get_code_cell(f"update_nb_metadata('{Path(fn).name}'{fmt_params})", hidden=False)
+    metadata_nb['cells'].append(metadata_cell)
+    write_nb(metadata_nb, meta_fn)
+
+def update_nb_metadata(nb_path=None, title=None, summary=None, keywords='fastai', overwrite=True, **kwargs):
+    "Creates jekyll metadata for given notebook path."
+    nb = read_nb(nb_path)
+    data = {'title': title, 'summary': summary, 'keywords': keywords, **kwargs}
+    data = {k:v for (k,v) in data.items() if v is not None} # remove none values
+    if not data: return
+    nb['metadata']['jekyll'] = data
+    write_nb(nb, nb_path)
+    NotebookNotary().sign(nb)
+
+def has_metadata_cell(cells, fn):
+    for c in cells:
+        if re.search(f"update_nb_metadata\('{fn}'", c['source']): return c
+
+def stringify(s): return f'\'{s}\'' if isinstance(s, str) else s
+
+IMPORT_RE = re.compile(r"from (fastai[\.\w_]*)")
+def get_imported_modules(cells, nb_module_name=''):
+    "Finds all submodules of notebook - sorted by submodules > top level modules > manual imports. This gives notebook imports priority"
+    module_names = get_top_level_modules()
+    nb_imports = [match.group(1) for cell in cells for match in IMPORT_RE.finditer(cell['source']) if cell['cell_type'] == 'code']
+    parts = nb_module_name.split('.')
+    parent_modules = ['.'.join(parts[:(x+1)]) for x in range_of(parts)] # Imports parent modules - a.b.c = [a, a.b, a.b.c]
+    all_modules = module_names + nb_imports + parent_modules
+    mods = [import_mod(m, ignore_errors=True) for m in all_modules]
+    return [m for m in mods if m is not None]
+
+def get_top_level_modules(num_levels=1):
+    mod_dir = Path(import_mod('fastai').__file__).parent
+    filtered_n = filter(lambda x: x.count('.')<=num_levels, get_module_names(mod_dir))
+    return sorted(filtered_n, key=lambda s: s.count('.'), reverse=True) # Submodules first (sorted by periods)
+
+NEW_FT_HEADER = '## New Methods - Please document or move to the undocumented section'
+UNDOC_HEADER = '## Undocumented Methods - Methods moved below this line will intentionally be hidden'
+def parse_sections(cells):
+    old_cells, undoc_cells, new_cells = [], [], []
+    current_section = old_cells
+    for cell in cells:
+        if cell['cell_type'] == 'markdown':
+            if re.match(UNDOC_HEADER, cell['source']): current_section = undoc_cells
+            if re.match(NEW_FT_HEADER, cell['source']): current_section = new_cells
+        current_section.append(cell)
+    undoc_cells = undoc_cells or [get_md_cell(UNDOC_HEADER)]
+    new_cells = new_cells or [get_md_cell(NEW_FT_HEADER)]
+    return old_cells, undoc_cells, new_cells
+
+def remove_undoc_cells(cells):
+    old, _, _ = parse_sections(cells)
+    return old
+
+# currently code vbox sub-cells mainly
+def remove_code_cell_jupyter_widget_state_elem(cells):
+    for c in cells:
+        if c['cell_type'] == 'code':
+            if 'outputs' in c:
+                c['outputs'] = [l for l in c['outputs'] if not ('data' in l and 'application/vnd.jupyter.widget-view+json' in l.data)]
+    return cells
+
+def update_module_page(mod, dest_path='.'):
+    "Update the documentation notebook of a given module."
+    doc_path = get_doc_path(mod, dest_path)
+    strip_name = strip_fastai(mod.__name__)
+    nb = read_nb(doc_path)
+    cells = nb['cells']
+
+    link_markdown_cells(cells, get_imported_modules(cells, mod.__name__))
+
+    type_dict = read_nb_types(cells)
+    gvar_map = get_global_vars(mod)
+    for name in get_exports(mod):
+        if name not in gvar_map: continue
+        code = gvar_map[name]
+        if name in type_dict: cells[type_dict[name]] = get_md_cell(code)
+        else: cells.append(get_md_cell(code))
+
+    pos_dict = read_nb_content(cells, strip_name)
+    ft_names = get_ft_names(mod, include_inner=True)
+    new_fts = list(set(ft_names) - set(pos_dict.keys()))
+    if new_fts: print(f'Found new fuctions for {mod}. Please document:\n{new_fts}')
+    existing, undoc_cells, new_cells = parse_sections(cells)
+    for ft_name in new_fts: new_cells.extend([get_doc_cell(ft_name), get_empty_cell()])
+    if len(new_cells) > 1: nb['cells'] = existing + undoc_cells + new_cells
+
+    write_nb(nb, doc_path)
+    return doc_path
+
+def link_nb(nb_path):
+    nb = read_nb(nb_path)
+    cells = nb['cells']
+    link_markdown_cells(cells, get_imported_modules(cells, Path(nb_path).stem))
+    write_nb(nb, nb_path)
+    NotebookNotary().sign(read_nb(nb_path))
+
+def get_module_from_notebook(doc_path):
+    "Find module given a source path. Assume it belongs to fastai directory"
+    return f'fastai.{Path(doc_path).stem}'
+
+def check_nbconvert_version():
+    import nbconvert
+    assert nbconvert.version_info >= (5,4,0), "Please update nbconvert to >=5.4 for consistent .html output"
+
+def update_notebooks(source_path, dest_path=None, update_html=True, document_new_fns=False,
+                     update_nb_links=True, html_path=None, force=False):
+    "`source_path` can be a directory or a file. Assume all modules reside in the fastai directory."
+    from .convert2html import convert_nb
+    source_path = Path(source_path)
+
+    if source_path.is_file():
+        dest_path = source_path.parent if dest_path is None else Path(dest_path)
+        html_path = dest_path/'..'/'docs' if html_path is None else Path(html_path)
+        doc_path = source_path
+        assert source_path.suffix == '.ipynb', 'Must update from notebook or module'
+        if document_new_fns:
+            mod = import_mod(get_module_from_notebook(source_path))
+            if not mod: print('Could not find module for path:', source_path)
+            elif mod.__file__.endswith('__init__.py'): pass
+            else: update_module_page(mod, dest_path)
+        generate_missing_metadata(doc_path)
+        if update_nb_links:
+            print(f'Updating notebook {doc_path}. Please wait...')
+            link_nb(doc_path)
+            execute_nb(doc_path, {'metadata': {'path': doc_path.parent}}, show_doc_only=True)
+        if update_html:
+            check_nbconvert_version()
+            html_fn = html_path/doc_path.with_suffix('.html').name
+            if not force and html_fn.is_file():
+                in_mod  = os.path.getmtime(doc_path)
+                out_mod = os.path.getmtime(html_fn)
+                if in_mod < out_mod: return
+            convert_nb(doc_path, html_path)
+
+    elif (source_path.name.startswith('fastai.')):
+        # Do module update
+        assert dest_path is not None, 'To update a module, you must specify a destination folder for where notebook resides'
+        mod = import_mod(source_path.name)
+        if not mod: return print('Could not find module for:', source_path)
+        doc_path = Path(dest_path)/(strip_fastai(mod.__name__)+'.ipynb')
+        if not doc_path.exists():
+            print('Notebook does not exist. Creating:', doc_path)
+            create_module_page(mod, dest_path)
+        update_notebooks(doc_path, dest_path=dest_path, update_html=update_html, document_new_fns=document_new_fns,
+                         update_nb_links=update_nb_links, html_path=html_path)
+    elif source_path.is_dir():
+        for f in sorted(Path(source_path).glob('*.ipynb')):
+            update_notebooks(f, dest_path=dest_path, update_html=update_html, document_new_fns=document_new_fns,
+                             update_nb_links=update_nb_links, html_path=html_path)
+    else: print('Could not resolve source file:', source_path)