Basic functions using pytorch

Torch Core

This module contains all the basic functions we need in other modules of the fastai library (split with core that contains the ones not requiring pytorch). Its documentation can easily be skipped at a first read, unless you want to know what a given fuction does.

Global constants

AdamW = partial(optim.Adam, betas=(0.9,0.99))

bn_types = (nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d)

default_device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')

Functions that operate conversions

flatten[source]

flatten(m)

Flattens all the layers of m.

model2half[source]

model2half(model:Module) → Module

Convert model to half precision except the batchnorm layers.

requires_grad[source]

requires_grad(m:Module, b:Optional[bool]=None) → Optional[bool]

If b is None, returns the requires_grad state of the first layer of m. Otherwise, sets requires_grad=b in all children of m.

tensor[source]

tensor(x:Any, rest) → Tensor

Like torch.as_tensor, but handle lists too, and can pass multiple vector elements directly

Ensures x is a torch Tensor.

to_data[source]

to_data(b:ItemsList)

Recursively map lists of items in b to their wrapped data

to_detach[source]

to_detach(b:Tensors)

Recursively detach lists of tensors in b

to_device[source]

to_device(b:Tensors, device:device)

Ensure b is on device.

to_half[source]

to_half(b:Collection[Tensor]) → Collection[Tensor]

Put the input of the batch b in half precision.

to_np[source]

to_np(x)

Convert x to a numpy array.

Functions to deal with model initialization

apply_init[source]

apply_init(m, init_func:LayerFunc)

Initialize all non-batchnorm layers of m with init_func.

apply_leaf[source]

apply_leaf(m:Module, f:LayerFunc)

Apply f to children of m.

cond_init[source]

cond_init(m:Module, init_func:LayerFunc)

Initialize the non-batchnorm layers of m with init_func

in_channels[source]

in_channels(m:Module) → List[int]

Return the shape of the first weight layer in m.

Functions to get information of a model

children[source]

children(m:Module) → ModuleList

Get children of module m.

first_layer[source]

first_layer(m:Module) → Module

Retrieve first layer in a module m.

num_children[source]

num_children(m:Module) → int

Get number of children modules in module m.

range_children[source]

range_children(m:Module) → Iterator[int]

Return iterator of len of children of m.

trainable_params[source]

trainable_params(m:Module) → ParamList

Return list of trainable params in m.

Functions to deal with BatchNorm layers

bn2float[source]

bn2float(module:Module) → Module

If module is batchnorm don't use half precision.

set_bn_eval[source]

set_bn_eval(m:Module)

Set bn layers in eval mode for all recursive children of m.

split_bn_bias[source]

split_bn_bias(layer_groups:ModuleList) → ModuleList

Sort each layer in layer_groups into batchnorm (bn_types) and non-batchnorm groups.

Other functions

calc_loss[source]

calc_loss(y_pred:Tensor, y_true:Tensor, loss_func:LossFunction)

Calculate loss between y_pred and y_true using loss_class and bs.

data_collate[source]

data_collate(batch:ItemsList) → Tensor

Convert batch items to tensor data.

split_model[source]

split_model(model:Module, splits:Collection[Union[Module, ModuleList]], want_idxs:bool=False)

Splits the model according to the layer in splits. If splits are layers, the model is split at those (not included) sequentially. If want_idxs is True, the corresponding indexes are returned. If splits are lists of layers, the model is split according to those.

split_model_idx[source]

split_model_idx(model:Module, idxs:Collection[int]) → ModuleList

Split model according to the indices in idxs.