How to use the batchflow.models.torch.TorchModel function in batchflow

"""  Ronneberger O. et al "`U-Net: Convolutional Networks for Biomedical Image Segmentation
`_"
"""
import numpy as np
import torch
import torch.nn as nn

from .layers import ConvBlock
from . import TorchModel
from .utils import get_shape


class UNet(TorchModel):
    """ UNet

    **Configuration**

    inputs : dict
        dict with 'images' and 'masks' (see :meth:`~.TorchModel._make_inputs`)

    body : dict
        num_blocks : int
            number of downsampling/upsampling blocks (default=4)

        filters : list of int
            number of filters in each block (default=[128, 256, 512, 1024])

        downsample : dict
            parameters for downsampling block

def default_config(cls):
        config = TorchModel.default_config()

        config['common'] += {'conv/bias': False}
        config['body/num_blocks'] = 5
        config['body/filters'] = (2 ** np.arange(config['body/num_blocks']) * 64).tolist()
        config['body/downsample'] = dict(layout='p', pool_size=2, pool_strides=2)
        config['body/encoder'] = dict(layout='cna cna', kernel_size=3)
        config['body/upsample'] = dict(layout='tna', kernel_size=2, strides=2)
        config['body/decoder'] = dict(layout='cna cna', kernel_size=3)
        config['head'] += dict(layout='c', kernel_size=1)

        config['loss'] = 'ce'

        return config

def default_config(cls):
        config = TorchModel.default_config()
        config['common/conv/bias'] = False
        config['initial_block'] = dict(layout='cnap', filters=64, kernel_size=7, strides=2,
                                       pool_size=3, pool_strides=2)

        config['body/block'] = dict(layout=None, post_activation=None, downsample=False,
                                    bottleneck=False, bottleneck_factor=4,
                                    width_factor=1, zero_pad=False,
                                    resnext=False, resnext_factor=32,
                                    se_block=False, se_factor=16)

        config['head'] = dict(layout='Vdf', dropout_rate=.4)

        config['loss'] = 'ce'
        # The learning rate starts from 0.1 (no warming up), and is divided by 10 at 30 and 60 epochs
        # with batch size = 256 on ImageNet.
        #config['decay'] = ('const', dict(boundaries=[117188, 234375], values=[lr, lr/10, lr/100]))

Sergey Zagoruyko, Nikos Komodakis. "`Wide Residual Networks
`_"

Xie S. et al. "`Aggregated Residual Transformations for Deep Neural Networks
`_"
"""
import numpy as np
import torch.nn as nn
import torch.nn.functional as F

from .layers import ConvBlock
from . import TorchModel
from .utils import get_shape, get_num_dims, get_num_channels


class ResNet(TorchModel):
    """ The base ResNet model

    Notes
    -----
    This class is intended to define custom ResNets.
    For more convenience use predefined :class:`~.torch.ResNet18`, :class:`~.torch.ResNet34`,
    and others described down below.

    **Configuration**

    inputs : dict
        dict with 'images' and 'labels' (see :meth:`~.TorchModel._make_inputs`)

    initial_block : dict
        filters : int
            number of filters (default=64)

""" Iandola F. et al. "`SqueezeNet: AlexNet-level accuracy with 50x fewer parameters and <0.5MB model size"
`_"
"""
import numpy as np
import torch
import torch.nn as nn

from . import TorchModel
from .layers import ConvBlock
from .utils import get_num_channels, get_shape


class SqueezeNet(TorchModel):
    """ SqueezeNet neural network

    **Configuration**

    inputs : dict
        dict with keys 'images' and 'labels' (see :meth:`~.TorchModel._make_inputs`)

    body : dict
        layout : str
            A sequence of blocks:

            - f : fire block
            - m : max-pooling
            - b : bypass

            Default is 'fffmffffmf'.

"""  Ronneberger O. et al "`U-Net: Convolutional Networks for Biomedical Image Segmentation
`_"
"""
import numpy as np
import torch
import torch.nn as nn

from .layers import ConvBlock
from . import TorchModel
from .utils import get_shape


class UNet(TorchModel):
    """ UNet

    **Configuration**

    inputs : dict
        dict with 'images' and 'masks' (see :meth:`~.TorchModel._make_inputs`)

    body : dict
        num_blocks : int
            number of downsampling/upsampling blocks (default=4)

        filters : list of int
            number of filters in each block (default=[128, 256, 512, 1024])

        downsample : dict
            parameters for downsampling block

_VGG19_ARCH = [
    (2, 0, 64, 1),
    (2, 0, 128, 1),
    (4, 0, 256, 1),
    (4, 0, 512, 1),
    (4, 0, 512, 1)
]

_VGG7_ARCH = [
    (2, 0, 64, 1),
    (2, 0, 128, 1),
    (2, 1, 256, 1)
]


class VGG(TorchModel):
    """ Base VGG neural network

    **Configuration**

    inputs : dict
        dict with keys 'images' and 'labels' (see :meth:`~.TorchModel._make_inputs`)

    body/arch : list of tuple of int
        Each list item contains parameters for one network block as a tuple of 4 ints:

        - number of convolution layers with 3x3 kernel
        - number of convolution layers with 1x1 kernel
        - number of filters in each layer
        - whether to downscale the image at the end of the block with max_pooling (2x2, stride=2)
    """
    @classmethod

`_"

Xie S. et al. "`Aggregated Residual Transformations for Deep Neural Networks
`_"
"""
import numpy as np
import torch.nn as nn
import torch.nn.functional as F

from ... import is_best_practice
from .layers import ConvBlock
from . import TorchModel
from .utils import get_shape, get_num_dims, get_num_channels


class ResNet(TorchModel):
    """ The base ResNet model

    Notes
    -----
    This class is intended to define custom ResNets.
    For more convenience use predefined :class:`~.torch.ResNet18`, :class:`~.torch.ResNet34`,
    and others described down below.

    **Configuration**

    inputs : dict
        dict with 'images' and 'labels' (see :meth:`~.TorchModel._make_inputs`)

    initial_block : dict
        filters : int
            number of filters (default=64)

"""  Encoder-decoder """
import torch.nn as nn

from .layers import ConvBlock
from . import TorchModel
from .resnet import ResNet18


class EncoderDecoder(TorchModel):
    """ Encoder-decoder architecture

    **Configuration**

    inputs : dict
        dict with 'images' (see :meth:`~.TorchModel._make_inputs`)

    body : dict
        encoder : dict
            base_class : TorchModel
                a model implementing ``make_encoder`` method which returns tensors
                with encoded representation of the inputs
            other args
                parameters for base class ``make_encoder`` method

        embedding : dict