model¶
- trojanzoo.utils.model.init_weights(m, filter_list=[])[source]¶
Traverse module
m
to intialize weights of all submodules except for those infilter_list
.Module parameters are reset by calling
module.reset_parameters()
.Note
An alternative implementation is to call
m.apply(init_weights)
so that this method could be non-recursive and avoid traverse.- Parameters:
m (torch.nn.Module) – Module to initialize.
filter_list (tuple[type]) – List of submodule types as exceptions. Defaults to
[]
(empty).
- Example:
from trojanzoo.utils.model import init_weights import torch.nn as nn net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2)) init_weights(filter_list=[nn.Linear]) # no change init_weights(net) # init nn.Linear layers
- trojanzoo.utils.model.get_layer_name(module, depth=-1, prefix='', use_filter=True, non_leaf=False, seq_only=False, init=True)[source]¶
Get layer names of a
torch.nn.Module
.- Parameters:
module (torch.nn.Module) – the module to process.
depth (int) – The traverse depth. Defaults to
-1
().prefix (str) – The prefix string to all elements. Defaults to empty string
''
.use_filter (bool) –
Whether to filter out certain layer types.
non_leaf (bool) – Whether to include non-leaf nodes. Defaults to
False
.seq_only (bool) – Whether to only traverse children of
torch.nn.Sequential
. IfFalse
, will traverse children of alltorch.nn.Module
. Defaults toFalse
.
- Returns:
list[str] – The list of all layer names.
- Example:
>>> import torchvision >>> from trojanzoo.utils.model import get_layer_name >>> >>> model = torchvision.models.resnet18() >>> get_layer_name(model, depth=0) [] >>> get_layer_name(model, depth=1) ['conv1', 'maxpool', 'layer1', 'layer2', 'layer3', 'layer4', 'avgpool', 'fc'] >>> get_layer_name(model, depth=2, prefix='model') ['model.conv1', 'model.maxpool', 'model.layer1.0', 'model.layer1.1', 'model.layer2.0', 'model.layer2.1', 'model.layer3.0', 'model.layer3.1', 'model.layer4.0', 'model.layer4.1', 'model.avgpool', 'model.fc'] >>> get_layer_name(model, seq_only=True) ['conv1', 'maxpool', 'layer1.0', 'layer1.1', 'layer2.0', 'layer2.1', 'layer3.0', 'layer3.1', 'layer4.0', 'layer4.1', 'avgpool', 'fc'] >>> get_layer_name(model, seq_only=True, non_leaf=True) ['conv1', 'maxpool', 'layer1.0', 'layer1.1', 'layer1', 'layer2.0', 'layer2.1', 'layer2', 'layer3.0', 'layer3.1', 'layer3', 'layer4.0', 'layer4.1', 'layer4', 'avgpool', 'fc'] >>> get_layer_name(model) ['conv1', 'maxpool', 'layer1.0.conv1', 'layer1.0.conv2', 'layer1.1.conv1', 'layer1.1.conv2', 'layer2.0.conv1', 'layer2.0.conv2', 'layer2.0.downsample.0', 'layer2.1.conv1', 'layer2.1.conv2', 'layer3.0.conv1', 'layer3.0.conv2', 'layer3.0.downsample.0', 'layer3.1.conv1', 'layer3.1.conv2', 'layer4.0.conv1', 'layer4.0.conv2', 'layer4.0.downsample.0', 'layer4.1.conv1', 'layer4.1.conv2', 'avgpool', 'fc']
- trojanzoo.utils.model.get_all_layer(module, x, layer_input='input', depth=-1, prefix='', use_filter=True, non_leaf=False, seq_only=True, verbose=0)[source]¶
Get all intermediate layer outputs of
_input
from any intermediate layer in atorch.nn.Module
.- Parameters:
module (torch.nn.Module) – the module to process.
x (torch.Tensor) – The batched input tensor from
layer_input
.layer_input (str) – The intermediate layer name of
x
. Defaults to'input'
.depth (int) – The traverse depth. Defaults to
-1
().prefix (str) – The prefix string to all elements. Defaults to empty string
''
.use_filter (bool) –
Whether to filter out certain layer types.
non_leaf (bool) – Whether to include non-leaf nodes. Defaults to
False
.seq_only (bool) – Whether to only traverse children of
torch.nn.Sequential
. IfFalse
, will traverse children of alltorch.nn.Module
. Defaults toFalse
.verbose (int) –
The output level to show information including layer name, output shape and module information. Setting it larger than
0
will enable the output. Different integer values stands for different module information. Defaults to0
.0
: No output1
: Show layer class name.2
: Show layer string (first line).3
: Show layer string (full).
- Returns:
dict[str, torch.Tensor] – The dict of all layer outputs.
- Example:
>>> import torch >>> import torchvision >>> from trojanzoo.utils.model import get_all_layer >>> >>> model = torchvision.models.densenet121() >>> x = torch.randn(5, 3, 224, 224) >>> y = get_all_layer(model.features, x, verbose=True) layer name output shape module information conv0 [5, 64, 112, 112] Conv2d pool0 [5, 64, 56, 56] MaxPool2d denseblock1 [5, 256, 56, 56] _DenseBlock transition1.conv [5, 128, 56, 56] Conv2d transition1.pool [5, 128, 28, 28] AvgPool2d denseblock2 [5, 512, 28, 28] _DenseBlock transition2.conv [5, 256, 28, 28] Conv2d transition2.pool [5, 256, 14, 14] AvgPool2d denseblock3 [5, 1024, 14, 14] _DenseBlock transition3.conv [5, 512, 14, 14] Conv2d transition3.pool [5, 512, 7, 7] AvgPool2d denseblock4 [5, 1024, 7, 7] _DenseBlock >>> y.keys() dict_keys(['conv0', 'pool0', 'denseblock1', 'transition1.conv', 'transition1.pool', 'denseblock2', 'transition2.conv', 'transition2.pool', 'denseblock3', 'transition3.conv', 'transition3.pool', 'denseblock4'])
Note
This method regards
module
as atorch.nn.Sequential
. Many modules embed flatten operation in theirforward
method (e.g.,view(n, -1)
orflatten(1)
), makingget_all_layer
raise error. We suggest to usetorch.nn.Flatten
instead to keep the module sequential.
- trojanzoo.utils.model.get_layer(module, x, layer_output='output', layer_input='input', layer_name_list=None, seq_only=True)[source]¶
Get one certain intermediate layer output of
_input
from any intermediate layer in atorch.nn.Module
.- Parameters:
module (torch.nn.Module) – the module to process.
x (torch.Tensor) – The batched input tensor from
layer_input
.layer_output (str) – The intermediate output layer name. Defaults to
'classifier'
.layer_input (str) – The intermediate layer name that outputs
x
. Defaults to'input'
.seq_only (bool) – Whether to only traverse children of
torch.nn.Sequential
. IfFalse
, will traverse children of alltorch.nn.Module
. Defaults toTrue
.
- Returns:
torch.Tensor – The output of layer
layer_output
.- Example:
>>> import torch >>> import torchvision >>> from trojanzoo.utils.model import get_all_layer, get_layer >>> >>> model = torchvision.models.densenet121() >>> x = torch.randn(5, 3, 224, 224) >>> y = get_all_layer(model.features, x, verbose=True) layer name output shape module information conv0 [5, 64, 112, 112] Conv2d pool0 [5, 64, 56, 56] MaxPool2d denseblock1 [5, 256, 56, 56] _DenseBlock transition1.conv [5, 128, 56, 56] Conv2d transition1.pool [5, 128, 28, 28] AvgPool2d denseblock2 [5, 512, 28, 28] _DenseBlock transition2.conv [5, 256, 28, 28] Conv2d transition2.pool [5, 256, 14, 14] AvgPool2d denseblock3 [5, 1024, 14, 14] _DenseBlock transition3.conv [5, 512, 14, 14] Conv2d transition3.pool [5, 512, 7, 7] AvgPool2d denseblock4 [5, 1024, 7, 7] _DenseBlock >>> x = torch.randn(6, 256, 56, 56) >>> get_layer(model.features, x, layer_input='denseblock1', layer_output='transition3.conv').shape torch.Size([6, 512, 14, 14])
- trojanzoo.utils.model.summary(module, depth=0, verbose=True, indent=0, tree_length=None, indent_atom=12)[source]¶
- Prints a string summary of the module.This method is similar to tensorflow.keras.Model.summary().
- Parameters:
module (torch.nn.Module) – The module to process.
depth (int) – The traverse depth. Defaults to
0
.verbose (bool) – Whether to output auxiliary information. Defaults to
True
.indent (int) – The space indent for the entire string. Defaults to
0
.tree_length (int) – The tree length. If
None
, useindent_atom * (depth + 1)
. Defaults toNone
.indent_atom (int) – The indent incremental for each sub-structure. Defaults to
12
.
- Example:
>>> import torchvision >>> from trojanzoo.utils.model import summary >>> >>> model=torchvision.models.resnet18() >>> summary(model) >>> summary(model, depth=1) conv1 Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False) bn1 BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) relu ReLU(inplace=True) maxpool MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False) layer1 Sequential layer2 Sequential layer3 Sequential layer4 Sequential avgpool AdaptiveAvgPool2d(output_size=(1, 1)) fc Linear(in_features=512, out_features=1000, bias=True) >>> summary(model, depth=1, verbose=False) conv1 bn1 relu maxpool layer1 layer2 layer3 layer4 avgpool fc >>> summary(model, depth=2) conv1 Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False) bn1 BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) relu ReLU(inplace=True) maxpool MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False) layer1 Sequential 0 BasicBlock 1 BasicBlock layer2 Sequential 0 BasicBlock 1 BasicBlock layer3 Sequential 0 BasicBlock 1 BasicBlock layer4 Sequential 0 BasicBlock 1 BasicBlock avgpool AdaptiveAvgPool2d(output_size=(1, 1)) fc Linear(in_features=512, out_features=1000, bias=True)
Note
You could use
get_all_layer()
withverbose=True
to see the output tensor shape for each layer.
- trojanzoo.utils.model.activate_params(module, params=[])[source]¶
Set
requires_grad=True
for selectedparams
ofmodule
. All other params are frozen.- Parameters:
module (torch.nn.Module) – The module to process.
params (Iterator[torch.nn.parameter.Parameter]) –
- The parameters to
requires_grad
. Defaults to
[]
.
- The parameters to
- trojanzoo.utils.model.accuracy(_output, _label, num_classes, topk=(1, 5))[source]¶
Computes the accuracy over the k top predictions for the specified values of k.
- Parameters:
_output (torch.Tensor) – The batched logit tensor with shape
(N, C)
._label (torch.Tensor) – The batched label tensor with shape
(N)
.num_classes (int) – Number of classes.
topk (Iterable[int]) – Which top-k accuracies to show. Defaults to
(1, 5)
.
- Returns:
list[float] – Top-k accuracies.
- trojanzoo.utils.model.generate_target(module, _input, idx=1, same=False)[source]¶
- Generate target labels of a batched input based on
the classification confidence ranking index.
- Parameters:
module (torch.nn.Module) – The module to process.
_input (torch.Tensor) – The input tensor.
idx (int) – The classification confidence rank of target class. Defaults to
1
.same (bool) – Generate the same label for all samples using mod. Defaults to
False
.
- Returns:
torch.Tensor – The generated target label with shape
(N)
.
- class trojanzoo.utils.model.ExponentialMovingAverage(model, decay)[source]¶
Maintains moving averages of model parameters using an exponential decay.
ema_avg = decay * avg_model_param + (1 - decay) * model_param
See also
https://github.com/pytorch/vision/blob/main/references/classification/utils.py
torch.optim.swa_utils.AveragedModel is used to compute the EMA.