ares.model package

class ares.model.cifar10_cls.CifarCLS(model_name, normalize=True)[source]

Bases: Module

The class to create cifar10 model.

__init__(model_name, normalize=True)[source]
Parameters:
  • model_name (str) – The model name in the cifar10 model zoo.

  • normalize (bool) – Whether interating the normalization layer into the model.

forward(x)[source]
Parameters:

x (torch.Tensor) – The input images. The images should be torch.Tensor with shape [N, C, H, W] and range [0, 1].

Returns:

The output logits with shape [N D].

Return type:

torch.Tensor

load()[source]

The function to load ckpt.

ares.model.cifar10_cls.filter_state_dict(state_dict)[source]

The function to filter state dict

class ares.model.imagenet_cls.ImageNetCLS(model_name, normalize=True)[source]

Bases: Module

The class to create ImageNet model.

__init__(model_name, normalize=True)[source]
Parameters:
  • model_name (str) – The model name in the ImageNet model zoo.

  • normalize (bool) – Whether interating the normalization layer into the model.

forward(x)[source]
Parameters:

x (torch.Tensor) – The input images. The images should be torch.Tensor with shape [N, C, H, W] and range [0, 1].

Returns:

The output logits with shape [N D].

Return type:

torch.Tensor

load()[source]

The function to load ckpt.

class ares.model.resnet.ResNet(block, layers, num_classes=1000, zero_init_residual=False, groups=1, width_per_group=64, replace_stride_with_dilation=None, norm_layer=None)[source]

Bases: Module

__init__(block, layers, num_classes=1000, zero_init_residual=False, groups=1, width_per_group=64, replace_stride_with_dilation=None, norm_layer=None)[source]

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(x)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

ares.model.resnet.resnet101(pretrained=False, progress=True, **kwargs)[source]

ResNet-101 model from “Deep Residual Learning for Image Recognition” <https://arxiv.org/pdf/1512.03385.pdf>

Parameters:
  • pretrained (bool) – If True, returns a model pre-trained on ImageNet

  • progress (bool) – If True, displays a progress bar of the download to stderr

ares.model.resnet.resnet152(pretrained=False, progress=True, **kwargs)[source]

ResNet-152 model from “Deep Residual Learning for Image Recognition” <https://arxiv.org/pdf/1512.03385.pdf>

Parameters:
  • pretrained (bool) – If True, returns a model pre-trained on ImageNet

  • progress (bool) – If True, displays a progress bar of the download to stderr

ares.model.resnet.resnet18(pretrained=False, progress=True, **kwargs)[source]

ResNet-18 model from “Deep Residual Learning for Image Recognition” <https://arxiv.org/pdf/1512.03385.pdf>

Parameters:
  • pretrained (bool) – If True, returns a model pre-trained on ImageNet

  • progress (bool) – If True, displays a progress bar of the download to stderr

ares.model.resnet.resnet34(pretrained=False, progress=True, **kwargs)[source]

ResNet-34 model from “Deep Residual Learning for Image Recognition” <https://arxiv.org/pdf/1512.03385.pdf>

Parameters:
  • pretrained (bool) – If True, returns a model pre-trained on ImageNet

  • progress (bool) – If True, displays a progress bar of the download to stderr

ares.model.resnet.resnet50(pretrained=False, progress=True, **kwargs)[source]

ResNet-50 model from “Deep Residual Learning for Image Recognition” <https://arxiv.org/pdf/1512.03385.pdf>

Parameters:
  • pretrained (bool) – If True, returns a model pre-trained on ImageNet

  • progress (bool) – If True, displays a progress bar of the download to stderr

ares.model.resnet.resnext101_32x8d(pretrained=False, progress=True, **kwargs)[source]

ResNeXt-101 32x8d model from “Aggregated Residual Transformation for Deep Neural Networks” <https://arxiv.org/pdf/1611.05431.pdf>

Parameters:
  • pretrained (bool) – If True, returns a model pre-trained on ImageNet

  • progress (bool) – If True, displays a progress bar of the download to stderr

ares.model.resnet.resnext50_32x4d(pretrained=False, progress=True, **kwargs)[source]

ResNeXt-50 32x4d model from “Aggregated Residual Transformation for Deep Neural Networks” <https://arxiv.org/pdf/1611.05431.pdf>

Parameters:
  • pretrained (bool) – If True, returns a model pre-trained on ImageNet

  • progress (bool) – If True, displays a progress bar of the download to stderr

ares.model.resnet.wide_resnet101_2(pretrained=False, progress=True, **kwargs)[source]

Wide ResNet-101-2 model from “Wide Residual Networks” <https://arxiv.org/pdf/1605.07146.pdf> The model is the same as ResNet except for the bottleneck number of channels which is twice larger in every block. The number of channels in outer 1x1 convolutions is the same, e.g. last block in ResNet-50 has 2048-512-2048 channels, and in Wide ResNet-50-2 has 2048-1024-2048.

Parameters:
  • pretrained (bool) – If True, returns a model pre-trained on ImageNet

  • progress (bool) – If True, displays a progress bar of the download to stderr

ares.model.resnet.wide_resnet50_2(pretrained=False, progress=True, **kwargs)[source]

Wide ResNet-50-2 model from “Wide Residual Networks” <https://arxiv.org/pdf/1605.07146.pdf> The model is the same as ResNet except for the bottleneck number of channels which is twice larger in every block. The number of channels in outer 1x1 convolutions is the same, e.g. last block in ResNet-50 has 2048-512-2048 channels, and in Wide ResNet-50-2 has 2048-1024-2048.

Parameters:
  • pretrained (bool) – If True, returns a model pre-trained on ImageNet

  • progress (bool) – If True, displays a progress bar of the download to stderr

ares.model.resnet.wide_resnet50_3(pretrained=False, progress=True, **kwargs)[source]

Wide ResNet-50-3 model

Parameters:
  • pretrained (bool) – If True, returns a model pre-trained on ImageNet

  • progress (bool) – If True, displays a progress bar of the download to stderr

ares.model.resnet.wide_resnet50_4(pretrained=False, progress=True, **kwargs)[source]

Wide ResNet-50-4 model

Parameters:
  • pretrained (bool) – If True, returns a model pre-trained on ImageNet

  • progress (bool) – If True, displays a progress bar of the download to stderr

ares.model.resnet.wide_resnet50_5(pretrained=False, progress=True, **kwargs)[source]

Wide ResNet-50-5 model

Parameters:
  • pretrained (bool) – If True, returns a model pre-trained on ImageNet

  • progress (bool) – If True, displays a progress bar of the download to stderr

ares.model.resnet.wide_resnet50_6(pretrained=False, progress=True, **kwargs)[source]

Wide ResNet-50-6 model

Parameters:
  • pretrained (bool) – If True, returns a model pre-trained on ImageNet

  • progress (bool) – If True, displays a progress bar of the download to stderr

class ares.model.preact_resnet.PreActBlock(in_planes, planes, stride=1)[source]

Bases: Module

Pre-activation version of the BasicBlock.

__init__(in_planes, planes, stride=1)[source]

Initializes internal Module state, shared by both nn.Module and ScriptModule.

expansion = 1
forward(x)[source]

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class ares.model.preact_resnet.PreActBottleneck(in_planes, planes, stride=1)[source]

Bases: Module

Pre-activation version of the original Bottleneck module.

__init__(in_planes, planes, stride=1)[source]

Initializes internal Module state, shared by both nn.Module and ScriptModule.

expansion = 4
forward(x)[source]

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class ares.model.preact_resnet.PreActResNet(block, num_blocks, num_classes=10)[source]

Bases: Module

The PreActResNet class.

__init__(block, num_blocks, num_classes=10)[source]

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(x)[source]

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

ares.model.preact_resnet.create_preact_res18()[source]

The function to create PreActResNet model.

class ares.model.wideresnet.BasicBlock(in_planes, out_planes, stride, dropRate=0.0)[source]

Bases: Module

__init__(in_planes, out_planes, stride, dropRate=0.0)[source]

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(x)[source]

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class ares.model.wideresnet.NetworkBlock(nb_layers, in_planes, out_planes, block, stride, dropRate=0.0)[source]

Bases: Module

__init__(nb_layers, in_planes, out_planes, block, stride, dropRate=0.0)[source]

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(x)[source]

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class ares.model.wideresnet.WideResNet(depth=34, num_classes=10, widen_factor=1, dropRate=0.0, use_FNandWN=False)[source]

Bases: Module

__init__(depth=34, num_classes=10, widen_factor=1, dropRate=0.0, use_FNandWN=False)[source]

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(x)[source]

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

ares.model.wideresnet.create_wres28_10()[source]

The function to create wide-resnet28-10 for cifar10 models.

ares.model.wideresnet.create_wres34_10()[source]

The function to create wide-resnet34-10 for cifar10 models.

ares.model.wideresnet.create_wres34_10_fn()[source]

The function to create wide-resnet34-10 with FN and WN for cifar10 models.

ares.model.resnet_denoise.resnet152_fd(**kwargs)[source]

The function to create resnet152 model of feature denoising.

class ares.model.vit_mae.VisionTransformer(global_pool_mae=True, **kwargs)[source]

Bases: VisionTransformer

Vision Transformer with support for global average pooling

__init__(global_pool_mae=True, **kwargs)[source]
Parameters:
  • img_size (int, tuple) – input image size

  • patch_size (int, tuple) – patch size

  • in_chans (int) – number of input channels

  • num_classes (int) – number of classes for classification head

  • global_pool (str) – type of global pooling for final sequence (default: ‘token’)

  • embed_dim (int) – embedding dimension

  • depth (int) – depth of transformer

  • num_heads (int) – number of attention heads

  • mlp_ratio (int) – ratio of mlp hidden dim to embedding dim

  • qkv_bias (bool) – enable bias for qkv if True

  • init_values – (float): layer-scale init values

  • class_token (bool) – use class token

  • fc_norm (Optional[bool]) – pre-fc norm after pool, set if global_pool == ‘avg’ if None (default: None)

  • drop_rate (float) – dropout rate

  • attn_drop_rate (float) – attention dropout rate

  • drop_path_rate (float) – stochastic depth rate

  • weight_init (str) – weight init scheme

  • embed_layer (nn.Module) – patch embedding layer

  • norm_layer – (nn.Module): normalization layer

  • act_layer – (nn.Module): MLP activation layer

forward_features(x)[source]
ares.model.vit_mae.vit_base_patch16(**kwargs)[source]

The function to create vit_base_patch16 model in MAE.

ares.model.vit_mae.vit_large_patch16(**kwargs)[source]

The function to create vit_large_patch16 model in MAE.