# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/callback.channelslast.ipynb (unless otherwise specified).


from __future__ import annotations


__all__ = ['ChannelsLastTfm', 'ChannelsLastCallback']

# Cell
#nbdev_comment from __future__ import annotations

from torch.cuda.amp import GradScaler

from fastcore.transform import DisplayedTransform, _is_tuple

from fastai.torch_core import TensorImageBase, TensorMask
from fastai.learner import Learner
from fastai.callback.core import Callback
from fastai.callback.fp16 import MixedPrecision
from fastai.basics import Pipeline

from ..imports import *

# Cell
class ChannelsLastTfm(DisplayedTransform):
    "Sets image-like inputs to `channels_last` format. For use in ChannelsLastCallback"
    order = 110 # run after all other transforms if added to batch_tfms
    def encodes(self, x:TensorImageBase):
        return x.to(memory_format=torch.channels_last)

# Cell
class ChannelsLastCallback(Callback):
    "Channels last training using PyTorch's Channels Last Memory Format (beta)"
    order = 8 # Before CastToTensor
    def __init__(self):
        self._channels_last = Pipeline([ChannelsLastTfm()])

    def before_fit(self):
        self.learn.model.to(memory_format=torch.channels_last)

    def before_batch(self):
        self.learn.xb = self._channels_last(self.xb)

# Cell
@patch
@delegates(GradScaler)
def to_channelslast(self:Learner, to_fp16=True, **kwargs):
    "Set `Learner` and inputs to `channels_last` format and Mixed Precision by default"
    if to_fp16 and not hasattr(self, 'mixed_precision') and not hasattr(self, 'channels_last'):
        return self.add_cbs([ChannelsLastCallback(), MixedPrecision(**kwargs)])
    elif not hasattr(self, 'channels_last'):
        return self.add_cb(ChannelsLastCallback())

# Cell
@patch
def to_contiguous(self:Learner, to_fp32=False):
    "Set `Learner` and inputs to `contiguous_format` (default format), optionally to single precision"
    self.model.to(memory_format=torch.contiguous_format)
    if to_fp32: return self.remove_cbs([ChannelsLastCallback, MixedPrecision])
    else:       return self.remove_cb(ChannelsLastCallback)