, bn_1st:bool=True, act_cls:type=ReLU, transpose:bool=False,
init:str='auto', xtra:NoneType=None, bias_std:float=0.01, dilation:Union=1, groups:int=1,
padding_mode:Literal='zeros', device:NoneType=None, dtype:NoneType=None
):
```
*Create a sequence of convolutional (`ni` to `nf`), ReLU (if
`use_activ`) and `norm_type` layers.*
The convolution uses `ks` (kernel size) `stride`, `padding` and `bias`.
`padding` will default to the appropriate value (`(ks-1)//2` if it’s not
a transposed conv) and `bias` will default to `True` the `norm_type` is
`Spectral` or `Weight`, `False` if it’s `Batch` or `BatchZero`. Note
that if you don’t want any normalization, you should pass
`norm_type=None`.
This defines a conv layer with `ndim` (1,2 or 3) that will be a
ConvTranspose if `transpose=True`. `act_cls` is the class of the
activation function to use (instantiated inside). Pass `act=None` if you
don’t want an activation function. If you quickly want to change your
default activation, you can change the value of `defaults.activation`.
`init` is used to initialize the weights (the bias are initialized to 0)
and `xtra` is an optional layer to add at the end.
``` python
tst = ConvLayer(16, 32)
mods = list(tst.children())
test_eq(len(mods), 3)
test_eq(mods[1].weight, torch.ones(32))
test_eq(mods[0].padding, (1,1))
```
``` python
x = torch.randn(64, 16, 8, 8)#.cuda()
```
``` python
#Padding is selected to make the shape the same if stride=1
test_eq(tst(x).shape, [64,32,8,8])
```
``` python
#Padding is selected to make the shape half if stride=2
tst = ConvLayer(16, 32, stride=2)
test_eq(tst(x).shape, [64,32,4,4])
```
``` python
#But you can always pass your own padding if you want
tst = ConvLayer(16, 32, padding=0)
test_eq(tst(x).shape, [64,32,6,6])
```
``` python
#No bias by default for Batch NormType
assert mods[0].bias is None
#But can be overridden with `bias=True`
tst = ConvLayer(16, 32, bias=True)
assert first(tst.children()).bias is not None
#For no norm, or spectral/weight, bias is True by default
for t in [None, NormType.Spectral, NormType.Weight]:
tst = ConvLayer(16, 32, norm_type=t)
assert first(tst.children()).bias is not None
```
``` python
#Various n_dim/tranpose
tst = ConvLayer(16, 32, ndim=3)
assert isinstance(list(tst.children())[0], nn.Conv3d)
tst = ConvLayer(16, 32, ndim=1, transpose=True)
assert isinstance(list(tst.children())[0], nn.ConvTranspose1d)
```
``` python
#No activation/leaky
tst = ConvLayer(16, 32, ndim=3, act_cls=None)
mods = list(tst.children())
test_eq(len(mods), 2)
tst = ConvLayer(16, 32, ndim=3, act_cls=partial(nn.LeakyReLU, negative_slope=0.1))
mods = list(tst.children())
test_eq(len(mods), 3)
assert isinstance(mods[2], nn.LeakyReLU)
```
``` python
# #export
# def linear(in_features, out_features, bias=True, act_cls=None, init='auto'):
# "Linear layer followed by optional activation, with optional auto-init"
# res = nn.Linear(in_features, out_features, bias=bias)
# if act_cls: act_cls = act_cls()
# init_linear(res, act_cls, init=init)
# if act_cls: res = nn.Sequential(res, act_cls)
# return res
```
``` python
# #export
# @delegates(ConvLayer)
# def conv1d(ni, nf, ks, stride=1, ndim=1, norm_type=None, **kwargs):
# "Convolutional layer followed by optional activation, with optional auto-init"
# return ConvLayer(ni, nf, ks, stride=stride, ndim=ndim, norm_type=norm_type, **kwargs)
```
``` python
# #export
# @delegates(ConvLayer)
# def conv2d(ni, nf, ks, stride=1, ndim=2, norm_type=None, **kwargs):
# "Convolutional layer followed by optional activation, with optional auto-init"
# return ConvLayer(ni, nf, ks, stride=stride, ndim=ndim, norm_type=norm_type, **kwargs)
```
``` python
# #export
# @delegates(ConvLayer)
# def conv3d(ni, nf, ks, stride=1, ndim=3, norm_type=None, **kwargs):
# "Convolutional layer followed by optional activation, with optional auto-init"
# return ConvLayer(ni, nf, ks, stride=stride, ndim=ndim, norm_type=norm_type, **kwargs)
```
------------------------------------------------------------------------
source
### AdaptiveAvgPool
``` python
def AdaptiveAvgPool(
sz:int=1, ndim:int=2
):
```
*nn.AdaptiveAvgPool layer for `ndim`*
------------------------------------------------------------------------
source
### MaxPool
``` python
def MaxPool(
ks:int=2, stride:NoneType=None, padding:int=0, ndim:int=2, ceil_mode:bool=False
):
```
*nn.MaxPool layer for `ndim`*
------------------------------------------------------------------------
source
### AvgPool
``` python
def AvgPool(
ks:int=2, stride:NoneType=None, padding:int=0, ndim:int=2, ceil_mode:bool=False
):
```
*nn.AvgPool layer for `ndim`*
## Embeddings
------------------------------------------------------------------------
source
### trunc_normal\_
``` python
def trunc_normal_(
x, mean:float=0.0, std:float=1.0
):
```
*Truncated normal initialization (approximation)*
------------------------------------------------------------------------
source
### Embedding
``` python
def Embedding(
ni, nf, std:float=0.01
):
```
*Embedding layer with truncated normal initialization*
Truncated normal initialization bounds the distribution to avoid large
value. For a given standard deviation `std`, the bounds are roughly
`-2*std`, `2*std`.
``` python
std = 0.02
tst = Embedding(10, 30, std)
assert tst.weight.min() > -2*std
assert tst.weight.max() < 2*std
test_close(tst.weight.mean(), 0, 1e-2)
test_close(tst.weight.std(), std, 0.1)
```
## Self attention
------------------------------------------------------------------------
source
### SelfAttention
``` python
def SelfAttention(
n_channels
):
```
*Self attention layer for `n_channels`.*
Self-attention layer as introduced in [Self-Attention Generative
Adversarial Networks](https://arxiv.org/abs/1805.08318).
Initially, no change is done to the input. This is controlled by a
trainable parameter named `gamma` as we return `x + gamma * out`.
``` python
tst = SelfAttention(16)
x = torch.randn(32, 16, 8, 8)
test_eq(tst(x),x)
```
Then during training `gamma` will probably change since it’s a trainable
parameter. Let’s see what’s happening when it gets a nonzero value.
``` python
tst.gamma.data.fill_(1.)
y = tst(x)
test_eq(y.shape, [32,16,8,8])
```
The attention mechanism requires three matrix multiplications (here
represented by 1x1 convs). The multiplications are done on the channel
level (the second dimension in our tensor) and we flatten the feature
map (which is 8x8 here). As in the paper, we note `f`, `g` and `h` the
results of those multiplications.
``` python
q,k,v = tst.query[0].weight.data,tst.key[0].weight.data,tst.value[0].weight.data
test_eq([q.shape, k.shape, v.shape], [[2, 16, 1], [2, 16, 1], [16, 16, 1]])
f,g,h = map(lambda m: x.view(32, 16, 64).transpose(1,2) @ m.squeeze().t(), [q,k,v])
test_eq([f.shape, g.shape, h.shape], [[32,64,2], [32,64,2], [32,64,16]])
```
The key part of the attention layer is to compute attention weights for
each of our location in the feature map (here 8x8 = 64). Those are
positive numbers that sum to 1 and tell the model to pay attention to
this or that part of the picture. We make the product of `f` and the
transpose of `g` (to get something of size bs by 64 by 64) then apply a
softmax on the first dimension (to get the positive numbers that sum up
to 1). The result can then be multiplied with `h` transposed to get an
output of size bs by channels by 64, which we can then be viewed as an
output the same size as the original input.
The final result is then `x + gamma * out` as we saw before.
``` python
beta = F.softmax(torch.bmm(f, g.transpose(1,2)), dim=1)
test_eq(beta.shape, [32, 64, 64])
out = torch.bmm(h.transpose(1,2), beta)
test_eq(out.shape, [32, 16, 64])
test_close(y, x + out.view(32, 16, 8, 8), eps=1e-4)
```
------------------------------------------------------------------------
source
### PooledSelfAttention2d
``` python
def PooledSelfAttention2d(
n_channels
):
```
*Pooled self attention layer for 2d.*
Self-attention layer used in the [Big GAN
paper](https://arxiv.org/abs/1809.11096).
It uses the same attention as in
[`SelfAttention`](https://docs.fast.ai/layers.html#selfattention) but
adds a max pooling of stride 2 before computing the matrices `g` and
`h`: the attention is ported on one of the 2x2 max-pooled window, not
the whole feature map. There is also a final matrix product added at the
end to the output, before retuning `gamma * out + x`.
------------------------------------------------------------------------
source
### SimpleSelfAttention
``` python
def SimpleSelfAttention(
n_in:int, ks:int=1, sym:bool=False
):
```
*Same as `nn.Module`, but no need for subclasses to call
`super().__init__`*
## PixelShuffle
PixelShuffle introduced in [this
article](https://arxiv.org/pdf/1609.05158.pdf) to avoid checkerboard
artifacts when upsampling images. If we want an output with `ch_out`
filters, we use a convolution with `ch_out * (r**2)` filters, where `r`
is the upsampling factor. Then we reorganize those filters like in the
picture below:
------------------------------------------------------------------------
source
### icnr_init
``` python
def icnr_init(
x, scale:int=2, init:function=kaiming_normal_
):
```
*ICNR init of `x`, with `scale` and `init` function*
ICNR init was introduced in [this
article](https://arxiv.org/abs/1707.02937). It suggests to initialize
the convolution that will be used in PixelShuffle so that each of the
`r**2` channels get the same weight (so that in the picture above, the 9
colors in a 3 by 3 window are initially the same).
> **Note**
>
> This is done on the first dimension because PyTorch stores the weights
> of a convolutional layer in this format: `ch_out x ch_in x ks x ks`.
``` python
tst = torch.randn(16*4, 32, 1, 1)
tst = icnr_init(tst)
for i in range(0,16*4,4):
test_eq(tst[i],tst[i+1])
test_eq(tst[i],tst[i+2])
test_eq(tst[i],tst[i+3])
```
------------------------------------------------------------------------
source
### PixelShuffle_ICNR
``` python
def PixelShuffle_ICNR(
ni, nf:NoneType=None, scale:int=2, blur:bool=False, norm_type:NormType=, act_cls:type=ReLU
):
```
*Upsample by `scale` from `ni` filters to `nf` (default `ni`), using
`nn.PixelShuffle`.*
The convolutional layer is initialized with
[`icnr_init`](https://docs.fast.ai/layers.html#icnr_init) and passed
`act_cls` and `norm_type` (the default of weight normalization seemed to
be what’s best for super-resolution problems, in our experiments).
The `blur` option comes from [Super-Resolution using Convolutional
Neural Networks without Any Checkerboard
Artifacts](https://arxiv.org/abs/1806.02658) where the authors add a
little bit of blur to completely get rid of checkerboard artifacts.
``` python
psfl = PixelShuffle_ICNR(16)
x = torch.randn(64, 16, 8, 8)
y = psfl(x)
test_eq(y.shape, [64, 16, 16, 16])
#ICNR init makes every 2x2 window (stride 2) have the same elements
for i in range(0,16,2):
for j in range(0,16,2):
test_eq(y[:,:,i,j],y[:,:,i+1,j])
test_eq(y[:,:,i,j],y[:,:,i ,j+1])
test_eq(y[:,:,i,j],y[:,:,i+1,j+1])
```
``` python
psfl = PixelShuffle_ICNR(16, norm_type=None)
x = torch.randn(64, 16, 8, 8)
y = psfl(x)
test_eq(y.shape, [64, 16, 16, 16])
#ICNR init makes every 2x2 window (stride 2) have the same elements
for i in range(0,16,2):
for j in range(0,16,2):
test_eq(y[:,:,i,j],y[:,:,i+1,j])
test_eq(y[:,:,i,j],y[:,:,i ,j+1])
test_eq(y[:,:,i,j],y[:,:,i+1,j+1])
```
``` python
psfl = PixelShuffle_ICNR(16, norm_type=NormType.Spectral)
x = torch.randn(64, 16, 8, 8)
y = psfl(x)
test_eq(y.shape, [64, 16, 16, 16])
#ICNR init makes every 2x2 window (stride 2) have the same elements
for i in range(0,16,2):
for j in range(0,16,2):
test_eq(y[:,:,i,j],y[:,:,i+1,j])
test_eq(y[:,:,i,j],y[:,:,i ,j+1])
test_eq(y[:,:,i,j],y[:,:,i+1,j+1])
```
## Sequential extensions
------------------------------------------------------------------------
source
### sequential
``` python
def sequential(
args:VAR_POSITIONAL
):
```
*Create an `nn.Sequential`, wrapping items with
[`Lambda`](https://docs.fast.ai/layers.html#lambda) if needed*
------------------------------------------------------------------------
source
### SequentialEx
``` python
def SequentialEx(
layers:VAR_POSITIONAL
):
```
*Like `nn.Sequential`, but with ModuleList semantics, and can access
module input*
This is useful to write layers that require to remember the input (like
a resnet block) in a sequential way.
------------------------------------------------------------------------
source
### MergeLayer
``` python
def MergeLayer(
dense:bool=False
):
```
*Merge a shortcut with the result of the module by adding them or
concatenating them if `dense=True`.*
``` python
res_block = SequentialEx(ConvLayer(16, 16), ConvLayer(16,16))
res_block.append(MergeLayer()) # just to test append - normally it would be in init params
x = torch.randn(32, 16, 8, 8)
y = res_block(x)
test_eq(y.shape, [32, 16, 8, 8])
test_eq(y, x + res_block[1](res_block[0](x)))
```
``` python
x = TensorBase(torch.randn(32, 16, 8, 8))
y = res_block(x)
test_is(y.orig, None)
```
## Concat
Equivalent to keras.layers.Concatenate, it will concat the outputs of a
ModuleList over a given dimension (default the filter dimension)
------------------------------------------------------------------------
source
### Cat
``` python
def Cat(
layers, dim:int=1
):
```
*Concatenate layers outputs over a given dim*
``` python
layers = [ConvLayer(2,4), ConvLayer(2,4), ConvLayer(2,4)]
x = torch.rand(1,2,8,8)
cat = Cat(layers)
test_eq(cat(x).shape, [1,12,8,8])
test_eq(cat(x), torch.cat([l(x) for l in layers], dim=1))
```
## Ready-to-go models
------------------------------------------------------------------------
source
### SimpleCNN
``` python
def SimpleCNN(
filters, kernel_szs:NoneType=None, strides:NoneType=None, bn:bool=True
):
```
*Create a simple CNN with `filters`.*
The model is a succession of convolutional layers from
`(filters[0],filters[1])` to `(filters[n-2],filters[n-1])` (if `n` is
the length of the `filters` list) followed by a
[`PoolFlatten`](https://docs.fast.ai/layers.html#poolflatten).
`kernel_szs` and `strides` defaults to a list of 3s and a list of 2s. If
`bn=True` the convolutional layers are successions of
conv-relu-batchnorm, otherwise conv-relu.
``` python
tst = SimpleCNN([8,16,32])
mods = list(tst.children())
test_eq(len(mods), 3)
test_eq([[m[0].in_channels, m[0].out_channels] for m in mods[:2]], [[8,16], [16,32]])
```
Test kernel sizes
``` python
tst = SimpleCNN([8,16,32], kernel_szs=[1,3])
mods = list(tst.children())
test_eq([m[0].kernel_size for m in mods[:2]], [(1,1), (3,3)])
```
Test strides
``` python
tst = SimpleCNN([8,16,32], strides=[1,2])
mods = list(tst.children())
test_eq([m[0].stride for m in mods[:2]], [(1,1),(2,2)])
```
------------------------------------------------------------------------
source
### ProdLayer
``` python
def ProdLayer(
):
```
*Merge a shortcut with the result of the module by multiplying them.*
------------------------------------------------------------------------
source
### SEModule
``` python
def SEModule(
ch, reduction, act_cls:type=ReLU
):
```
*Call self as a function.*
------------------------------------------------------------------------
source
### ResBlock
``` python
def ResBlock(
expansion, ni, nf, stride:int=1, groups:int=1, reduction:NoneType=None, nh1:NoneType=None, nh2:NoneType=None,
dw:bool=False, g2:int=1, sa:bool=False, sym:bool=False, norm_type:NormType=,
act_cls:type=ReLU, ndim:int=2, ks:int=3, pool:function=AvgPool, pool_first:bool=True, padding:NoneType=None,
bias:NoneType=None, bn_1st:bool=True, transpose:bool=False, init:str='auto', xtra:NoneType=None,
bias_std:float=0.01, dilation:Union=1, padding_mode:Literal='zeros', device:NoneType=None, dtype:NoneType=None
):
```
*Resnet block from `ni` to `nh` with `stride`*
This is a resnet block (normal or bottleneck depending on `expansion`, 1
for the normal block and 4 for the traditional bottleneck) that
implements the tweaks from [Bag of Tricks for Image Classification with
Convolutional Neural Networks](https://arxiv.org/abs/1812.01187). In
particular, the last batchnorm layer (if that is the selected
`norm_type`) is initialized with a weight (or gamma) of zero to
facilitate the flow from the beginning to the end of the network. It
also implements optional [Squeeze and
Excitation](https://arxiv.org/abs/1709.01507) and grouped convs for
[ResNeXT](https://arxiv.org/abs/1611.05431) and similar models (use
`dw=True` for depthwise convs).
The `kwargs` are passed to
[`ConvLayer`](https://docs.fast.ai/layers.html#convlayer) along with
`norm_type`.
------------------------------------------------------------------------
source
### SEBlock
``` python
def SEBlock(
expansion, ni, nf, groups:int=1, reduction:int=16, stride:int=1, kwargs:VAR_KEYWORD
):
```
*Call self as a function.*
------------------------------------------------------------------------
source
### SEResNeXtBlock
``` python
def SEResNeXtBlock(
expansion, ni, nf, groups:int=32, reduction:int=16, stride:int=1, base_width:int=4, kwargs:VAR_KEYWORD
):
```
*Call self as a function.*
------------------------------------------------------------------------
source
### SeparableBlock
``` python
def SeparableBlock(
expansion, ni, nf, reduction:int=16, stride:int=1, base_width:int=4, kwargs:VAR_KEYWORD
):
```
*Call self as a function.*
## Time Distributed Layer
Equivalent to Keras
[`TimeDistributed`](https://docs.fast.ai/layers.html#timedistributed)
Layer, enables computing pytorch
[`Module`](https://docs.fast.ai/torch_core.html#module) over an axis.
``` python
bs, seq_len = 2, 5
x, y = torch.rand(bs,seq_len,3,2,2), torch.rand(bs,seq_len,3,2,2)
```
``` python
tconv = TimeDistributed(nn.Conv2d(3,4,1))
test_eq(tconv(x).shape, (2,5,4,2,2))
tconv.low_mem=True
test_eq(tconv(x).shape, (2,5,4,2,2))
```
``` python
class Mod(Module):
def __init__(self):
self.conv = nn.Conv2d(3,4,1)
def forward(self, x, y):
return self.conv(x) + self.conv(y)
tmod = TimeDistributed(Mod())
```
``` python
out = tmod(x,y)
test_eq(out.shape, (2,5,4,2,2))
tmod.low_mem=True
out_low_mem = tmod(x,y)
test_eq(out_low_mem.shape, (2,5,4,2,2))
test_eq(out, out_low_mem)
```
``` python
class Mod2(Module):
def __init__(self):
self.conv = nn.Conv2d(3,4,1)
def forward(self, x, y):
return self.conv(x), self.conv(y)
tmod2 = TimeDistributed(Mod2())
```
``` python
out = tmod2(x,y)
test_eq(len(out), 2)
test_eq(out[0].shape, (2,5,4,2,2))
tmod2.low_mem=True
out_low_mem = tmod2(x,y)
test_eq(out_low_mem[0].shape, (2,5,4,2,2))
test_eq(out, out_low_mem)
```
------------------------------------------------------------------------
source
### TimeDistributed
``` python
def TimeDistributed(
module, low_mem:bool=False, tdim:int=1
):
```
*Applies [`module`](https://docs.fast.ai/layers.html#module) over `tdim`
identically for each step, use `low_mem` to compute one at a time.*
This module is equivalent to [Keras TimeDistributed
Layer](https://keras.io/api/layers/recurrent_layers/time_distributed/).
This wrapper allows to apply a layer to every temporal slice of an
input. By default it is assumed the time axis (`tdim`) is the 1st one
(the one after the batch size). A typical usage would be to encode a
sequence of images using an image encoder.
The `forward` function of
[`TimeDistributed`](https://docs.fast.ai/layers.html#timedistributed)
supports `*args` and `**kkwargs` but only `args` will be split and
passed to the underlying module independently for each timestep,
`kwargs` will be passed as they are. This is useful when you have module
that take multiple arguments as inputs, this way, you can put all
tensors you need spliting as `args` and other arguments that don’t need
split as `kwargs`.
> This module is heavy on memory, as it will try to pass mutiple
> timesteps at the same time on the batch dimension, if you get out of
> memorey errors, try first reducing your batch size by the number of
> timesteps.
``` python
from fastai.vision.all import *
```
``` python
encoder = create_body(resnet18())
```
A resnet18 will encode a feature map of 512 channels. Height and Width
will be divided by 32.
``` python
time_resnet = TimeDistributed(encoder)
```
a synthetic batch of 2 image-sequences of lenght 5.
`(bs, seq_len, ch, w, h)`
``` python
image_sequence = torch.rand(2, 5, 3, 64, 64)
```
``` python
time_resnet(image_sequence).shape
```
torch.Size([2, 5, 512, 2, 2])
This way, one can encode a sequence of images on feature space. There is
also a `low_mem_forward` that will pass images one at a time to reduce
GPU memory consumption.
``` python
time_resnet.low_mem_forward(image_sequence).shape
```
torch.Size([2, 5, 512, 2, 2])
## Swish and Mish
------------------------------------------------------------------------
source
### swish
``` python
def swish(
x, inplace:bool=False
):
```
*Call self as a function.*
------------------------------------------------------------------------
source
### SwishJit
``` python
def SwishJit(
):
```
*Same as `nn.Module`, but no need for subclasses to call
`super().__init__`*
------------------------------------------------------------------------
source
### MishJitAutoFn
``` python
def MishJitAutoFn(
args:VAR_POSITIONAL, kwargs:VAR_KEYWORD
):
```
*Base class to create custom `autograd.Function`.*
To create a custom `autograd.Function`, subclass this class and
implement the :meth:`forward` and :meth:`backward` static methods. Then,
to use your custom op in the forward pass, call the class method
`[`apply`](https://docs.fast.ai/torch_core.html#apply)`. Do not call
:meth:`forward` directly.
To ensure correctness and best performance, make sure you are calling
the correct methods on `ctx` and validating your backward function using
:func:`torch.autograd.gradcheck`.
See :ref:`extending-autograd` for more details on how to use this class.
Examples::
>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_AUTOGRAD)
>>> class Exp(Function):
>>> @staticmethod
>>> def forward(ctx, i):
>>> result = i.exp()
>>> ctx.save_for_backward(result)
>>> return result
>>>
>>> @staticmethod
>>> def backward(ctx, grad_output):
>>> result, = ctx.saved_tensors
>>> return grad_output * result
>>>
>>> # Use it by calling the apply method:
>>> # xdoctest: +SKIP
>>> output = Exp.apply(input)
------------------------------------------------------------------------
source
### mish
``` python
def mish(
x, inplace:bool=False
):
```
*Call self as a function.*
------------------------------------------------------------------------
source
### MishJit
``` python
def MishJit(
):
```
*Same as `nn.Module`, but no need for subclasses to call
`super().__init__`*
## Helper functions for submodules
It’s easy to get the list of all parameters of a given model. For when
you want all submodules (like linear/conv layers) without forgetting
lone parameters, the following class wraps those in fake modules.
------------------------------------------------------------------------
source
### ParameterModule
``` python
def ParameterModule(
p
):
```
*Register a lone parameter `p` in a module.*
------------------------------------------------------------------------
source
### children_and_parameters
``` python
def children_and_parameters(
m
):
```
*Return the children of `m` and its direct parameters not registered in
modules.*
``` python
class TstModule(Module):
def __init__(self): self.a,self.lin = nn.Parameter(torch.randn(1)),nn.Linear(5,10)
tst = TstModule()
children = children_and_parameters(tst)
test_eq(len(children), 2)
test_eq(children[0], tst.lin)
assert isinstance(children[1], ParameterModule)
test_eq(children[1].val, tst.a)
```
------------------------------------------------------------------------
source
### has_children
``` python
def has_children(
m
):
```
*Call self as a function.*
``` python
class A(Module): pass
assert not has_children(A())
assert has_children(TstModule())
```
------------------------------------------------------------------------
source
### flatten_model
``` python
def flatten_model(
m
):
```
*Return the list of all submodules and parameters of `m`*
``` python
tst = nn.Sequential(TstModule(), TstModule())
children = flatten_model(tst)
test_eq(len(children), 4)
assert isinstance(children[1], ParameterModule)
assert isinstance(children[3], ParameterModule)
```
------------------------------------------------------------------------
source
### NoneReduce
``` python
def NoneReduce(
loss_func
):
```
*A context manager to evaluate `loss_func` with none reduce.*
``` python
x,y = torch.randn(5),torch.randn(5)
loss_fn = nn.MSELoss()
with NoneReduce(loss_fn) as loss_func:
loss = loss_func(x,y)
test_eq(loss.shape, [5])
test_eq(loss_fn.reduction, 'mean')
loss_fn = F.mse_loss
with NoneReduce(loss_fn) as loss_func:
loss = loss_func(x,y)
test_eq(loss.shape, [5])
test_eq(loss_fn, F.mse_loss)
```
------------------------------------------------------------------------
source
### in_channels
``` python
def in_channels(
m
):
```
*Return the shape of the first weight layer in `m`.*
``` python
test_eq(in_channels(nn.Sequential(nn.Conv2d(5,4,3), nn.Conv2d(4,3,3))), 5)
test_eq(in_channels(nn.Sequential(nn.AvgPool2d(4), nn.Conv2d(4,3,3))), 4)
test_eq(in_channels(nn.Sequential(BatchNorm(4), nn.Conv2d(4,3,3))), 4)
test_eq(in_channels(nn.Sequential(InstanceNorm(4), nn.Conv2d(4,3,3))), 4)
test_eq(in_channels(nn.Sequential(InstanceNorm(4, affine=False), nn.Conv2d(4,3,3))), 4)
test_fail(lambda : in_channels(nn.Sequential(nn.AvgPool2d(4))))
```