diff --git a/configs/classification/imagenet/inception/inceptionv3_b32x8_100e.py b/configs/classification/imagenet/inception/inceptionv3_b32x8_100e.py index a0cf8d57..1d74ec39 100644 --- a/configs/classification/imagenet/inception/inceptionv3_b32x8_100e.py +++ b/configs/classification/imagenet/inception/inceptionv3_b32x8_100e.py @@ -4,16 +4,31 @@ # model settings model = dict( type='Classification', - backbone=dict(type='Inception3'), - head=dict( - type='ClsHead', - with_avg_pool=True, - in_channels=2048, - loss_config=dict( - type='CrossEntropyLossWithLabelSmooth', - label_smooth=0, + backbone=dict(type='Inception3', num_classes=1000), + head=[ + dict( + type='ClsHead', + with_fc=False, + in_channels=2048, + loss_config=dict( + type='CrossEntropyLossWithLabelSmooth', + label_smooth=0, + ), + num_classes=num_classes, + input_feature_index=[1], ), - num_classes=num_classes)) + dict( + type='ClsHead', + with_fc=False, + in_channels=768, + loss_config=dict( + type='CrossEntropyLossWithLabelSmooth', + label_smooth=0, + ), + num_classes=num_classes, + input_feature_index=[0], + ) + ]) class_list = [ '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', '10', '11', '12', '13', @@ -196,3 +211,5 @@ interval=10, hooks=[dict(type='TextLoggerHook'), dict(type='TensorboardLoggerHook')]) + +export = dict(export_type='raw', export_neck=True) diff --git a/configs/classification/imagenet/inception/inceptionv4_b32x8_100e.py b/configs/classification/imagenet/inception/inceptionv4_b32x8_100e.py new file mode 100644 index 00000000..09f00a3a --- /dev/null +++ b/configs/classification/imagenet/inception/inceptionv4_b32x8_100e.py @@ -0,0 +1,33 @@ +_base_ = 'configs/classification/imagenet/inception/inceptionv3_b32x8_100e.py' + +num_classes = 1000 +# model settings +model = dict( + type='Classification', + backbone=dict(type='Inception4', num_classes=num_classes), + head=[ + dict( + type='ClsHead', + with_fc=False, + in_channels=1536, + loss_config=dict( + type='CrossEntropyLossWithLabelSmooth', + label_smooth=0, + ), + num_classes=num_classes, + input_feature_index=[1], + ), + dict( + type='ClsHead', + with_fc=False, + in_channels=768, + loss_config=dict( + type='CrossEntropyLossWithLabelSmooth', + label_smooth=0, + ), + num_classes=num_classes, + input_feature_index=[0], + ) + ]) + +img_norm_cfg = dict(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) diff --git a/configs/classification/imagenet/inception/inceptionv4_b32x8_200e_rmsprop.py b/configs/classification/imagenet/inception/inceptionv4_b32x8_200e_rmsprop.py new file mode 100644 index 00000000..efa4a786 --- /dev/null +++ b/configs/classification/imagenet/inception/inceptionv4_b32x8_200e_rmsprop.py @@ -0,0 +1,46 @@ +# A config with the optimization settings from https://arxiv.org/pdf/1602.07261 +# May run with 20 GPUs +_base_ = 'configs/classification/imagenet/inception/inceptionv3_b32x8_100e.py' + +num_classes = 1000 +# model settings +model = dict( + type='Classification', + backbone=dict(type='Inception4', num_classes=num_classes), + head=[ + dict( + type='ClsHead', + with_fc=False, + in_channels=1536, + loss_config=dict( + type='CrossEntropyLossWithLabelSmooth', + label_smooth=0, + ), + num_classes=num_classes, + input_feature_index=[1], + ), + dict( + type='ClsHead', + with_fc=False, + in_channels=768, + loss_config=dict( + type='CrossEntropyLossWithLabelSmooth', + label_smooth=0, + ), + num_classes=num_classes, + input_feature_index=[0], + ) + ]) + +img_norm_cfg = dict(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) + +# optimizer +optimizer = dict( + type='RMSprop', lr=0.045, momentum=0.9, weight_decay=0.9, eps=1.0) + +# learning policy +lr_config = dict(policy='exp', gamma=0.96954) # gamma**2 ~ 0.94 +checkpoint_config = dict(interval=10) + +# runtime settings +total_epochs = 200 diff --git a/configs/classification/imagenet/mobilenet/mobilenetv2.py b/configs/classification/imagenet/mobilenet/mobilenetv2.py index 29a663f7..354d966a 100644 --- a/configs/classification/imagenet/mobilenet/mobilenetv2.py +++ b/configs/classification/imagenet/mobilenet/mobilenetv2.py @@ -13,7 +13,8 @@ type='CrossEntropyLossWithLabelSmooth', label_smooth=0, ), - num_classes=num_classes)) + num_classes=num_classes), + pretrained=True) # optimizer optimizer = dict(type='SGD', lr=0.1, momentum=0.9, weight_decay=0.0001) @@ -25,4 +26,4 @@ # runtime settings total_epochs = 100 checkpoint_sync_export = True -export = dict(export_neck=True) +export = dict(export_type='raw', export_neck=True) diff --git a/configs/classification/imagenet/resnext/resnext50-32x4d_b32x8_100e_jpg.py b/configs/classification/imagenet/resnext/resnext50-32x4d_b32x8_100e_jpg.py index c1156bd2..c5cacc07 100644 --- a/configs/classification/imagenet/resnext/resnext50-32x4d_b32x8_100e_jpg.py +++ b/configs/classification/imagenet/resnext/resnext50-32x4d_b32x8_100e_jpg.py @@ -19,7 +19,8 @@ type='CrossEntropyLossWithLabelSmooth', label_smooth=0, ), - num_classes=num_classes)) + num_classes=num_classes), + pretrained=True) # optimizer optimizer = dict(type='SGD', lr=0.1, momentum=0.9, weight_decay=0.0001) @@ -30,3 +31,4 @@ # runtime settings total_epochs = 100 +export = dict(export_type='raw', export_neck=True) diff --git a/easycv/apis/export.py b/easycv/apis/export.py index 85424a62..701d9f23 100644 --- a/easycv/apis/export.py +++ b/easycv/apis/export.py @@ -158,34 +158,49 @@ def _get_blade_model(): def _export_onnx_cls(model, model_config, cfg, filename, meta): + support_backbones = { + 'ResNet': { + 'depth': [50] + }, + 'MobileNetV2': {}, + 'Inception3': {}, + 'Inception4': {}, + 'ResNeXt': { + 'depth': [50] + } + } + if model_config['backbone'].get('type', None) not in support_backbones: + tmp = ' '.join(support_backbones.keys()) + info_str = f'Only support export onnx model for {tmp} now!' + raise ValueError(info_str) + configs = support_backbones[model_config['backbone'].get('type')] + for k, v in configs.items(): + if v[0].__class__(model_config['backbone'].get(k, None)) not in v: + raise ValueError( + f"Unsupport config for {model_config['backbone'].get('type')}") + + # save json config for test_pipline and class + with io.open( + filename + + '.config.json' if filename.endswith('onnx') else filename + + '.onnx.config.json', 'w') as ofile: + json.dump(meta, ofile) - if model_config['backbone'].get( - 'type', None) == 'ResNet' and model_config['backbone'].get( - 'depth', None) == 50: - # save json config for test_pipline and class - with io.open( - filename + - '.config.json' if filename.endswith('onnx') else filename + - '.onnx.config.json', 'w') as ofile: - json.dump(meta, ofile) - - device = 'cuda' if torch.cuda.is_available() else 'cpu' - model.eval() - model.to(device) - img_size = int(cfg.image_size2) - x_input = torch.randn((1, 3, img_size, img_size)).to(device) - torch.onnx.export( - model, - (x_input, 'onnx'), - filename if filename.endswith('onnx') else filename + '.onnx', - export_params=True, - opset_version=12, - do_constant_folding=True, - input_names=['input'], - output_names=['output'], - ) - else: - raise ValueError('Only support export onnx model for ResNet now!') + device = 'cuda' if torch.cuda.is_available() else 'cpu' + model.eval() + model.to(device) + img_size = int(cfg.image_size2) + x_input = torch.randn((1, 3, img_size, img_size)).to(device) + torch.onnx.export( + model, + (x_input, 'onnx'), + filename if filename.endswith('onnx') else filename + '.onnx', + export_params=True, + opset_version=12, + do_constant_folding=True, + input_names=['input'], + output_names=['output'], + ) def _export_cls(model, cfg, filename): diff --git a/easycv/models/backbones/__init__.py b/easycv/models/backbones/__init__.py index 5dc4561b..794f4426 100644 --- a/easycv/models/backbones/__init__.py +++ b/easycv/models/backbones/__init__.py @@ -10,6 +10,7 @@ from .genet import PlainNet from .hrnet import HRNet from .inceptionv3 import Inception3 +from .inceptionv4 import Inception4 from .lighthrnet import LiteHRNet from .mae_vit_transformer import * from .mit import MixVisionTransformer diff --git a/easycv/models/backbones/inceptionv3.py b/easycv/models/backbones/inceptionv3.py index 69f591bb..0439d9fa 100644 --- a/easycv/models/backbones/inceptionv3.py +++ b/easycv/models/backbones/inceptionv3.py @@ -2,9 +2,6 @@ r""" This model is taken from the official PyTorch model zoo. - torchvision.models.inception.py on 31th Aug, 2019 """ - -from collections import namedtuple - import torch import torch.nn as nn import torch.nn.functional as F @@ -16,8 +13,6 @@ __all__ = ['Inception3'] -_InceptionOutputs = namedtuple('InceptionOutputs', ['logits', 'aux_logits']) - @BACKBONES.register_module class Inception3(nn.Module): @@ -113,6 +108,7 @@ def forward(self, x): # N x 768 x 17 x 17 x = self.Mixed_6e(x) # N x 768 x 17 x 17 + aux = None if self.training and self.aux_logits: aux = self.AuxLogits(x) # N x 768 x 17 x 17 @@ -132,10 +128,7 @@ def forward(self, x): if hasattr(self, 'fc'): x = self.fc(x) - # N x 1000 (num_classes) - if self.training and self.aux_logits and hasattr(self, 'fc'): - return [_InceptionOutputs(x, aux)] - return [x] + return [aux, x] class InceptionA(nn.Module): diff --git a/easycv/models/backbones/inceptionv4.py b/easycv/models/backbones/inceptionv4.py new file mode 100644 index 00000000..917dbf0f --- /dev/null +++ b/easycv/models/backbones/inceptionv4.py @@ -0,0 +1,393 @@ +from __future__ import absolute_import, division, print_function +from collections import namedtuple + +import torch +import torch.nn as nn +import torch.nn.functional as F +from mmcv.cnn import constant_init, kaiming_init +from torch.nn.modules.batchnorm import _BatchNorm + +from ..modelzoo import inceptionv4 as model_urls +from ..registry import BACKBONES + +__all__ = ['Inception4'] + + +class BasicConv2d(nn.Module): + + def __init__(self, + in_planes, + out_planes, + kernel_size, + stride=1, + padding=0): + super(BasicConv2d, self).__init__() + self.conv = nn.Conv2d( + in_planes, + out_planes, + kernel_size=kernel_size, + stride=stride, + padding=padding, + bias=False) # verify bias false + self.bn = nn.BatchNorm2d( + out_planes, + eps=0.001, # value found in tensorflow + momentum=0.1, # default pytorch value + affine=True) + self.relu = nn.ReLU(inplace=True) + + def forward(self, x): + x = self.conv(x) + x = self.bn(x) + x = self.relu(x) + return x + + +class Mixed_3a(nn.Module): + + def __init__(self): + super(Mixed_3a, self).__init__() + self.maxpool = nn.MaxPool2d(3, stride=2) + self.conv = BasicConv2d(64, 96, kernel_size=3, stride=2) + + def forward(self, x): + x0 = self.maxpool(x) + x1 = self.conv(x) + out = torch.cat((x0, x1), 1) + return out + + +class Mixed_4a(nn.Module): + + def __init__(self): + super(Mixed_4a, self).__init__() + + self.branch0 = nn.Sequential( + BasicConv2d(160, 64, kernel_size=1, stride=1), + BasicConv2d(64, 96, kernel_size=3, stride=1)) + + self.branch1 = nn.Sequential( + BasicConv2d(160, 64, kernel_size=1, stride=1), + BasicConv2d(64, 64, kernel_size=(1, 7), stride=1, padding=(0, 3)), + BasicConv2d(64, 64, kernel_size=(7, 1), stride=1, padding=(3, 0)), + BasicConv2d(64, 96, kernel_size=(3, 3), stride=1)) + + def forward(self, x): + x0 = self.branch0(x) + x1 = self.branch1(x) + out = torch.cat((x0, x1), 1) + return out + + +class Mixed_5a(nn.Module): + + def __init__(self): + super(Mixed_5a, self).__init__() + self.conv = BasicConv2d(192, 192, kernel_size=3, stride=2) + self.maxpool = nn.MaxPool2d(3, stride=2) + + def forward(self, x): + x0 = self.conv(x) + x1 = self.maxpool(x) + out = torch.cat((x0, x1), 1) + return out + + +class Inception_A(nn.Module): + + def __init__(self): + super(Inception_A, self).__init__() + self.branch0 = BasicConv2d(384, 96, kernel_size=1, stride=1) + + self.branch1 = nn.Sequential( + BasicConv2d(384, 64, kernel_size=1, stride=1), + BasicConv2d(64, 96, kernel_size=3, stride=1, padding=1)) + + self.branch2 = nn.Sequential( + BasicConv2d(384, 64, kernel_size=1, stride=1), + BasicConv2d(64, 96, kernel_size=3, stride=1, padding=1), + BasicConv2d(96, 96, kernel_size=3, stride=1, padding=1)) + + self.branch3 = nn.Sequential( + nn.AvgPool2d(3, stride=1, padding=1, count_include_pad=False), + BasicConv2d(384, 96, kernel_size=1, stride=1)) + + def forward(self, x): + x0 = self.branch0(x) + x1 = self.branch1(x) + x2 = self.branch2(x) + x3 = self.branch3(x) + out = torch.cat((x0, x1, x2, x3), 1) + return out + + +class Reduction_A(nn.Module): + + def __init__(self): + super(Reduction_A, self).__init__() + self.branch0 = BasicConv2d(384, 384, kernel_size=3, stride=2) + + self.branch1 = nn.Sequential( + BasicConv2d(384, 192, kernel_size=1, stride=1), + BasicConv2d(192, 224, kernel_size=3, stride=1, padding=1), + BasicConv2d(224, 256, kernel_size=3, stride=2)) + + self.branch2 = nn.MaxPool2d(3, stride=2) + + def forward(self, x): + x0 = self.branch0(x) + x1 = self.branch1(x) + x2 = self.branch2(x) + out = torch.cat((x0, x1, x2), 1) + return out + + +class Inception_B(nn.Module): + + def __init__(self): + super(Inception_B, self).__init__() + self.branch0 = BasicConv2d(1024, 384, kernel_size=1, stride=1) + + self.branch1 = nn.Sequential( + BasicConv2d(1024, 192, kernel_size=1, stride=1), + BasicConv2d( + 192, 224, kernel_size=(1, 7), stride=1, padding=(0, 3)), + BasicConv2d( + 224, 256, kernel_size=(7, 1), stride=1, padding=(3, 0))) + + self.branch2 = nn.Sequential( + BasicConv2d(1024, 192, kernel_size=1, stride=1), + BasicConv2d( + 192, 192, kernel_size=(7, 1), stride=1, padding=(3, 0)), + BasicConv2d( + 192, 224, kernel_size=(1, 7), stride=1, padding=(0, 3)), + BasicConv2d( + 224, 224, kernel_size=(7, 1), stride=1, padding=(3, 0)), + BasicConv2d( + 224, 256, kernel_size=(1, 7), stride=1, padding=(0, 3))) + + self.branch3 = nn.Sequential( + nn.AvgPool2d(3, stride=1, padding=1, count_include_pad=False), + BasicConv2d(1024, 128, kernel_size=1, stride=1)) + + def forward(self, x): + x0 = self.branch0(x) + x1 = self.branch1(x) + x2 = self.branch2(x) + x3 = self.branch3(x) + out = torch.cat((x0, x1, x2, x3), 1) + return out + + +class Reduction_B(nn.Module): + + def __init__(self): + super(Reduction_B, self).__init__() + + self.branch0 = nn.Sequential( + BasicConv2d(1024, 192, kernel_size=1, stride=1), + BasicConv2d(192, 192, kernel_size=3, stride=2)) + + self.branch1 = nn.Sequential( + BasicConv2d(1024, 256, kernel_size=1, stride=1), + BasicConv2d( + 256, 256, kernel_size=(1, 7), stride=1, padding=(0, 3)), + BasicConv2d( + 256, 320, kernel_size=(7, 1), stride=1, padding=(3, 0)), + BasicConv2d(320, 320, kernel_size=3, stride=2)) + + self.branch2 = nn.MaxPool2d(3, stride=2) + + def forward(self, x): + x0 = self.branch0(x) + x1 = self.branch1(x) + x2 = self.branch2(x) + out = torch.cat((x0, x1, x2), 1) + return out + + +class Inception_C(nn.Module): + + def __init__(self): + super(Inception_C, self).__init__() + + self.branch0 = BasicConv2d(1536, 256, kernel_size=1, stride=1) + + self.branch1_0 = BasicConv2d(1536, 384, kernel_size=1, stride=1) + self.branch1_1a = BasicConv2d( + 384, 256, kernel_size=(1, 3), stride=1, padding=(0, 1)) + self.branch1_1b = BasicConv2d( + 384, 256, kernel_size=(3, 1), stride=1, padding=(1, 0)) + + self.branch2_0 = BasicConv2d(1536, 384, kernel_size=1, stride=1) + self.branch2_1 = BasicConv2d( + 384, 448, kernel_size=(3, 1), stride=1, padding=(1, 0)) + self.branch2_2 = BasicConv2d( + 448, 512, kernel_size=(1, 3), stride=1, padding=(0, 1)) + self.branch2_3a = BasicConv2d( + 512, 256, kernel_size=(1, 3), stride=1, padding=(0, 1)) + self.branch2_3b = BasicConv2d( + 512, 256, kernel_size=(3, 1), stride=1, padding=(1, 0)) + + self.branch3 = nn.Sequential( + nn.AvgPool2d(3, stride=1, padding=1, count_include_pad=False), + BasicConv2d(1536, 256, kernel_size=1, stride=1)) + + def forward(self, x): + x0 = self.branch0(x) + + x1_0 = self.branch1_0(x) + x1_1a = self.branch1_1a(x1_0) + x1_1b = self.branch1_1b(x1_0) + x1 = torch.cat((x1_1a, x1_1b), 1) + + x2_0 = self.branch2_0(x) + x2_1 = self.branch2_1(x2_0) + x2_2 = self.branch2_2(x2_1) + x2_3a = self.branch2_3a(x2_2) + x2_3b = self.branch2_3b(x2_2) + x2 = torch.cat((x2_3a, x2_3b), 1) + + x3 = self.branch3(x) + + out = torch.cat((x0, x1, x2, x3), 1) + return out + + +class InceptionAux(nn.Module): + + def __init__(self, in_channels, num_classes): + super(InceptionAux, self).__init__() + self.conv0 = BasicConv2d(in_channels, 128, kernel_size=1) + self.conv1 = BasicConv2d(128, 768, kernel_size=5) + self.conv1.stddev = 0.01 + self.fc = nn.Linear(768, num_classes) + self.fc.stddev = 0.001 + + def forward(self, x): + # N x 768 x 17 x 17 + x = F.avg_pool2d(x, kernel_size=5, stride=3) + # N x 768 x 5 x 5 + x = self.conv0(x) + # N x 128 x 5 x 5 + x = self.conv1(x) + # N x 768 x 1 x 1 + # Adaptive average pooling + x = F.adaptive_avg_pool2d(x, (1, 1)) + # N x 768 x 1 x 1 + x = torch.flatten(x, 1) + # N x 768 + x = self.fc(x) + # N x 1000 + return x + + +# class BasicConv2d(nn.Module): + +# def __init__(self, in_channels, out_channels, **kwargs): +# super(BasicConv2d, self).__init__() +# self.conv = nn.Conv2d(in_channels, out_channels, bias=False, **kwargs) +# self.bn = nn.BatchNorm2d(out_channels, eps=0.001) + +# def forward(self, x): +# x = self.conv(x) +# x = self.bn(x) +# return F.relu(x, inplace=True) + + +@BACKBONES.register_module +class Inception4(nn.Module): + """InceptionV4 backbone. + + Args: + num_classes (int): The num_classes of InceptionV4. An extra fc will be used if + """ + + def __init__(self, + num_classes: int = 0, + p_dropout=0.2, + aux_logits: bool = True): + super(Inception4, self).__init__() + self.aux_logits = aux_logits + # Modules + self.features = nn.Sequential( + BasicConv2d(3, 32, kernel_size=3, stride=2), + BasicConv2d(32, 32, kernel_size=3, stride=1), + BasicConv2d(32, 64, kernel_size=3, stride=1, padding=1), + Mixed_3a(), + Mixed_4a(), + Mixed_5a(), + Inception_A(), + Inception_A(), + Inception_A(), + Inception_A(), + Reduction_A(), # Mixed_6a + Inception_B(), + Inception_B(), + Inception_B(), + Inception_B(), + Inception_B(), + Inception_B(), + Inception_B(), # Mixed_6h 1024 x 17 x 17 + Reduction_B(), # Mixed_7a + Inception_C(), + Inception_C(), + Inception_C()) + + if aux_logits: + self.AuxLogits = InceptionAux(1024, num_classes) + + self.dropout = nn.Dropout(p_dropout) + self.last_linear = None + if num_classes > 0: + self.last_linear = nn.Linear(1536, num_classes) + + self.default_pretrained_model_path = model_urls[ + self.__class__.__name__] + + @property + def fc(self): + return self.last_linear + + def init_weights(self): + for m in self.modules(): + if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear): + for m in self.modules(): + if isinstance(m, nn.Conv2d): + kaiming_init(m, mode='fan_in', nonlinearity='relu') + elif isinstance(m, (_BatchNorm, nn.GroupNorm)): + constant_init(m, 1) + elif isinstance(m, nn.BatchNorm2d): + nn.init.constant_(m.weight, 1) + nn.init.constant_(m.bias, 0) + + def logits(self, features): + x = F.adaptive_avg_pool2d(features, output_size=(1, 1)) + # x = F.avg_pool2d(features, kernel_size=adaptiveAvgPoolWidth) + x = x.view(x.size(0), -1) # B x 1536 + x = self.fc(x) + # B x num_classes + return x + + def forward(self, input: torch.Tensor): + """_summary_ + + Args: + input (torch.Tensor): A RGB image tensor with shape B x C x H x W + + Returns: + torch.Tensor: A feature tensor or a logit tensor when num_classes is 0 (default) + """ + + if self.training and self.aux_logits: + x = self.features[:-4](input) + aux = self.AuxLogits(x) + x = self.features[-4:](x) + else: + x = self.features(input) + aux = None + + if self.fc is not None: + x = self.logits(x) + + return [aux, x] diff --git a/easycv/models/classification/classification.py b/easycv/models/classification/classification.py index 479b4d1e..3118c24a 100644 --- a/easycv/models/classification/classification.py +++ b/easycv/models/classification/classification.py @@ -151,6 +151,7 @@ def forward_backbone(self, img: torch.Tensor) -> List[torch.Tensor]: x = self.backbone(img) return x + @torch.jit.unused def forward_onnx(self, img: torch.Tensor) -> Dict[str, torch.Tensor]: """ forward_onnx means generate prob from image only support one neck + one head diff --git a/easycv/models/modelzoo.py b/easycv/models/modelzoo.py index d367cc62..baa49d37 100644 --- a/easycv/models/modelzoo.py +++ b/easycv/models/modelzoo.py @@ -79,6 +79,12 @@ 'http://pai-vision-data-hz.oss-cn-zhangjiakou.aliyuncs.com/pretrained_models/easycv/inceptionv3/inception_v3.pth', } +inceptionv4 = { + # Inception v4 ported from http://data.lip6.fr/cadene/pretrainedmodels/inceptionv4-8e4777a0.pth + 'Inception4': + 'http://pai-vision-data-hz.oss-cn-zhangjiakou.aliyuncs.com/pretrained_models/easycv/inceptionv4/inception_v4.pth', +} + genet = { 'PlainNetnormal': 'http://pai-vision-data-hz.oss-cn-zhangjiakou.aliyuncs.com/pretrained_models/easycv/genet/GENet_normal.pth', diff --git a/easycv/utils/config_tools.py b/easycv/utils/config_tools.py index e673f49a..631f50d0 100644 --- a/easycv/utils/config_tools.py +++ b/easycv/utils/config_tools.py @@ -515,6 +515,10 @@ def validate_export_config(cfg): 'configs/classification/imagenet/swint/imagenet_swin_tiny_patch4_window7_224_jpg.py', 'CLASSIFICATION_M0BILENET': 'configs/classification/imagenet/mobilenet/mobilenetv2.py', + 'CLASSIFICATION_INCEPTIONV4': + 'configs/classification/imagenet/inception/inceptionv4_b32x8_100e.py', + 'CLASSIFICATION_INCEPTIONV3': + 'configs/classification/imagenet/inception/inceptionv3_b32x8_100e.py', # metric learning 'METRICLEARNING': diff --git a/tests/test_models/backbones/test_inceptionv4.py b/tests/test_models/backbones/test_inceptionv4.py new file mode 100644 index 00000000..d9b7625c --- /dev/null +++ b/tests/test_models/backbones/test_inceptionv4.py @@ -0,0 +1,63 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +import copy +import random +import unittest + +import numpy as np +import torch + +from easycv.models import modelzoo +from easycv.models.backbones import Inception4 + + +class InceptionV3Test(unittest.TestCase): + + def setUp(self): + print(('Testing %s.%s' % (type(self).__name__, self._testMethodName))) + + def test_inceptionv3_withfc(self): + with torch.no_grad(): + # input data + batch_size = random.randint(10, 30) + a = torch.rand(batch_size, 3, 299, 299).to('cuda') + + num_classes = random.randint(10, 1000) + net = Inception4( + aux_logits=True, num_classes=num_classes).to('cuda') + net.init_weights() + net.train() + + self.assertTrue(len(list(net(a)[-1].shape)) == 2) + self.assertTrue(len(list(net(a)[0].shape)) == 2) + self.assertTrue(net(a)[-1].size(1) == num_classes) + self.assertTrue(net(a)[-1].size(0) == batch_size) + self.assertTrue(net(a)[0].size(1) == num_classes) + self.assertTrue(net(a)[0].size(0) == batch_size) + + def test_inceptionv3_withoutfc(self): + with torch.no_grad(): + # input data + batch_size = random.randint(10, 30) + a = torch.rand(batch_size, 3, 299, 299).to('cuda') + + net = Inception4(aux_logits=True, num_classes=0).to('cuda') + net.init_weights() + net.eval() + + self.assertTrue(net(a)[-1].size(1) == 1536) + self.assertTrue(net(a)[-1].size(0) == batch_size) + + def test_inceptionv3_load_modelzoo(self): + with torch.no_grad(): + net = Inception4(aux_logits=True, num_classes=1000).to('cuda') + original_weight = net.features[0].conv.weight + original_weight = copy.deepcopy(original_weight.cpu().data.numpy()) + + net.init_weights() + load_weight = net.features[0].conv.weight.cpu().data.numpy() + + self.assertFalse(np.allclose(original_weight, load_weight)) + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/test_tools/test_export.py b/tests/test_tools/test_export.py new file mode 100644 index 00000000..7b5a96cf --- /dev/null +++ b/tests/test_tools/test_export.py @@ -0,0 +1,138 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +import logging +import os +import sys +import unittest + +import numpy as np +import onnxruntime +import torch + +from easycv.models import build_model +from easycv.utils.checkpoint import load_checkpoint +from easycv.utils.config_tools import mmcv_config_fromfile, rebuild_config +from easycv.utils.test_util import run_in_subprocess + +sys.path.append(os.path.dirname(os.path.realpath(__file__))) +logging.basicConfig(level=logging.INFO) + +WORK_DIRECTORY = 'work_dir3' + +BASIC_EXPORT_CONFIGS = { + 'config_file': None, + 'checkpoint': 'dummy', + 'output_filename': f'{WORK_DIRECTORY}/test_out.pth', + 'user_config_params': ['--export.export_type', 'onnx'] +} + + +def build_cmd(export_configs, MODEL_TYPE) -> str: + base_cmd = 'python tools/export.py' + base_cmd += f" {export_configs['config_file']}" + base_cmd += f" {export_configs['checkpoint']}" + base_cmd += f" {export_configs['output_filename']}" + base_cmd += f' --model_type {MODEL_TYPE}' + user_params = ' '.join(export_configs['user_config_params']) + base_cmd += f' --user_config_params {user_params}' + return base_cmd + + +class ExportTest(unittest.TestCase): + """In this unittest, we test the onnx export functionality of + some classification/detection models. + """ + + def setUp(self): + print(('Testing %s.%s' % (type(self).__name__, self._testMethodName))) + os.makedirs(WORK_DIRECTORY, exist_ok=True) + + def tearDown(self): + super().tearDown() + + def run_test(self, + CONFIG_FILE, + MODEL_TYPE, + img_size: int = 224, + **override_configs): + configs = BASIC_EXPORT_CONFIGS.copy() + configs['config_file'] = CONFIG_FILE + + configs.update(override_configs) + + cmd = build_cmd(configs, MODEL_TYPE) + logging.info(f'Export with commands: {cmd}') + run_in_subprocess(cmd) + + cfg = mmcv_config_fromfile(configs['config_file']) + cfg = rebuild_config(cfg, configs['user_config_params']) + + if hasattr(cfg.model, 'pretrained'): + cfg.model.pretrained = False + + torch_model = build_model(cfg.model).eval() + if 'checkpoint' in override_configs: + load_checkpoint( + torch_model, + override_configs['checkpoint'], + strict=False, + logger=logging.getLogger()) + session = onnxruntime.InferenceSession(configs['output_filename'] + + '.onnx') + input_tensor = torch.randn((1, 3, img_size, img_size)) + + torch_output = torch_model(input_tensor, mode='test')['prob'] + + onnx_output = session.run( + [session.get_outputs()[0].name], + {session.get_inputs()[0].name: np.array(input_tensor)}) + if isinstance(onnx_output, list): + onnx_output = onnx_output[0] + + onnx_output = torch.tensor(onnx_output) + + is_same_shape = torch_output.shape == onnx_output.shape + + self.assertTrue( + is_same_shape, + f'The shapes of the two outputs are mismatch, got {torch_output.shape} and {onnx_output.shape}' + ) + is_allclose = torch.allclose(torch_output, onnx_output) + + torch_out_minmax = f'{float(torch_output.min())}~{float(torch_output.max())}' + onnx_out_minmax = f'{float(onnx_output.min())}~{float(onnx_output.max())}' + + info_msg = f'got avg: {float(torch_output.mean())} and {float(onnx_output.mean())},' + info_msg += f' and range: {torch_out_minmax} and {onnx_out_minmax}' + self.assertTrue( + is_allclose, + f'The values between the two outputs are mismatch, {info_msg}') + + def test_inceptionv3(self): + CONFIG_FILE = 'configs/classification/imagenet/inception/inceptionv3_b32x8_100e.py' + self.run_test(CONFIG_FILE, 'CLASSIFICATION_INCEPTIONV3', 299) + + def test_inceptionv4(self): + CONFIG_FILE = 'configs/classification/imagenet/inception/inceptionv4_b32x8_100e.py' + self.run_test(CONFIG_FILE, 'CLASSIFICATION_INCEPTIONV4', 299) + + def test_resnext50(self): + CONFIG_FILE = 'configs/classification/imagenet/resnext/imagenet_resnext50-32x4d_jpg.py' + self.run_test( + CONFIG_FILE, + 'CLASSIFICATION_RESNEXT', + checkpoint= + 'https://pai-vision-data-hz.oss-cn-zhangjiakou.aliyuncs.com/EasyCV/modelzoo/classification/resnext/resnext50-32x4d/epoch_100.pth' + ) + + def test_mobilenetv2(self): + CONFIG_FILE = 'configs/classification/imagenet/mobilenet/mobilenetv2.py' + self.run_test( + CONFIG_FILE, + 'CLASSIFICATION_M0BILENET', + checkpoint= + 'http://pai-vision-data-hz.oss-cn-zhangjiakou.aliyuncs.com/pretrained_models/easycv/mobilenetv2/mobilenet_v2.pth' + ) + + +if __name__ == '__main__': + unittest.main()