def __init__(self):
super(SSH, self).__init__()
# backbone
self.vgg16 = nn.ModuleList(make_layers(vgg_cfgs['D']))
# SSH - M3
self.M3 = M_Module(512, 256, 128)
self.M3_bbox_pred = nn.Conv2d(512, 8, 1, 1, 0)
self.M3_cls_score = nn.Conv2d(512, 4, 1, 1, 0)
self.M3_cls_score_softmax = nn.Softmax(dim=1)
# SSH - M2
self.M2 = M_Module(512, 256, 128)
self.M2_bbox_pred = nn.Conv2d(512, 8, 1, 1, 0)
self.M2_cls_score = nn.Conv2d(512, 4, 1, 1, 0)
self.M2_cls_score_softmax = nn.Softmax(dim=1)
# SSH - M1
self.conv4_128 = nn.Conv2d(512, 128, 1, 1, 0)
self.conv4_128_relu = nn.ReLU(inplace=True)
self.conv5_128 = nn.Conv2d(512, 128, 1, 1, 0)
self.conv5_128_relu = nn.ReLU(inplace=True)
self.conv5_128_up = nn.ConvTranspose2d(128, 128, 4, 2, 1, groups=128, bias=False)
self.eltadd = nn.EltAdd()
self.conv4_fuse_final = nn.Conv2d(128, 128, 3, 1, 1)
self.conv4_fuse_final_relu = nn.ReLU(inplace=True)
self.M1 = M_Module(128, 128, 64)
self.M1_bbox_pred = nn.Conv2d(256, 8, 1, 1, 0)
self.M1_cls_score = nn.Conv2d(256, 4, 1, 1, 0)
self.M1_cls_score_softmax = nn.Softmax(dim=1)
def __init__(self, backbone='ResNet50'):
super(RetinaFace, self).__init__()
if backbone.lower() == 'resnet50':
self.backbone = resnet50()
self.return_layers = {'layer2': 1, 'layer3': 2, 'layer4': 3}
self.in_channels_stage2 = 256
self.out_channels = 256
elif backbone.lower() == 'mobilenetv1':
self.backbone = MobileNetV1()
self.return_layers = {'stage1': 1, 'stage2': 2, 'stage3': 3}
self.in_channels_stage2 = 32
self.out_channels = 64
else:
raise NotImplementedError
self.body = _utils.IntermediateLayerGetter(self.backbone,
self.return_layers)
self.in_channels_list = [
self.in_channels_stage2 * 2, self.in_channels_stage2 * 4,
self.in_channels_stage2 * 8
]
self.fpn = FPN(self.in_channels_list, self.out_channels)
self.ssh1 = SSH(self.out_channels, self.out_channels)
self.ssh2 = SSH(self.out_channels, self.out_channels)
self.ssh3 = SSH(self.out_channels, self.out_channels)
self.ClassHead = self._make_class_head(fpn_num=3,
inchannels=self.out_channels)
self.BboxHead = self._make_bbox_head(fpn_num=3,
inchannels=self.out_channels)
self.LandmarkHead = self._make_landmark_head(
fpn_num=3, inchannels=self.out_channels)
self.softmax = nn.Softmax(dim=-1)
def _pyramidbox(self):
self.ssh_conv3_l2norm = nn.L2Norm2d(512, 10)
self.ssh_conv4_l2norm = nn.L2Norm2d(512, 8)
self.ssh_conv5_l2norm = nn.L2Norm2d(512, 5)
self.SSHchannels = [512, 512, 512, 512, 512, 512]
loc = []
conf = []
for i in range(6):
loc.append(
nn.Conv2d(self.SSHchannels[i],
8,
kernel_size=3,
stride=1,
padding=1))
if i == 0:
conf.append(
nn.Conv2d(self.SSHchannels[i],
8,
kernel_size=3,
stride=1,
padding=1))
else:
conf.append(
nn.Conv2d(self.SSHchannels[i],
6,
kernel_size=3,
stride=1,
padding=1))
self.mbox_loc = nn.ModuleList(loc)
self.mbox_conf = nn.ModuleList(conf)
self.softmax = nn.Softmax(dim=-1)
def __init__(self, model_type='32'):
super(EXTD, self).__init__()
self.model_type = int(model_type)
self.mobilenet = MobileNetV2(model_type=self.model_type)
if self.model_type == 32:
self.base = nn.ModuleList(self.mobilenet.features)[:8]
elif self.model_type == 48 or self.model_type == 64:
self.base = nn.ModuleList(self.mobilenet.features)[:6]
else:
raise NotImplementedError
self.upsample = nn.Upsample(scale_factor=2,
mode='bilinear',
align_corners=False)
self.eltadd = nn.EltAdd()
self.upfeat = []
for it in range(5):
self.upfeat.append(
upsample(in_channels=self.model_type,
out_channels=self.model_type))
self.upfeat = nn.ModuleList(self.upfeat)
self.loc = []
self.conf = []
self.net_source = [4, 4, 4, 4]
self.feature_dim = []
if self.model_type == 32:
self.feature_dim += [self.base[4].conv[-3].out_channels]
for idx in self.net_source:
self.feature_dim += [self.base[idx].conv[-3].out_channels]
else:
self.feature_dim += [self.base[4].conv[-2].out_channels]
for idx in self.net_source:
self.feature_dim += [self.base[idx].conv[-2].out_channels]
self.loc += [
nn.Conv2d(self.feature_dim[0], 4, kernel_size=3, padding=1)
]
self.conf += [
nn.Conv2d(self.feature_dim[0], 4, kernel_size=3, padding=1)
]
for k, v in enumerate(self.net_source, 1):
self.loc += [
nn.Conv2d(self.feature_dim[k], 4, kernel_size=3, padding=1)
]
self.conf += [
nn.Conv2d(self.feature_dim[k], 2, kernel_size=3, padding=1)
]
self.loc += [nn.Conv2d(self.model_type, 4, kernel_size=3, padding=1)]
self.conf += [nn.Conv2d(self.model_type, 2, kernel_size=3, padding=1)]
self.loc = nn.ModuleList(self.loc)
self.conf = nn.ModuleList(self.conf)
self.softmax = nn.Softmax(dim=-1)
def __init__(self):
super(SHF, self).__init__()
self.conv1 = self.build_conv_block(3, 64, with_pool=True)
self.conv2 = self.build_conv_block(64, 128, with_pool=True)
self.conv3 = self.build_conv_block(128, 256, n_conv=3, with_pool=True)
self.conv4 = self.build_conv_block(256, 512, n_conv=3)
self.pool4 = nn.MaxPool2d(2, 2)
self.conv5 = self.build_conv_block(512, 512, n_conv=3)
self.conv5_256 = self.build_conv_block(512,
256,
1,
padding=0,
n_conv=1)
self.conv5_256_up = nn.ConvTranspose2d(256,
256,
kernel_size=4,
stride=2,
padding=1,
groups=256,
bias=False)
self.conv4_256 = self.build_conv_block(512,
256,
1,
padding=0,
n_conv=1)
self.conv4_fuse_final = self.build_conv_block(512, 512, n_conv=1)
self.conv4_fuse_final_dim_red = self.build_conv_block(512,
128,
n_conv=1)
self.head1 = self.build_conv_block(128,
128,
padding=1,
dilation=1,
n_conv=1)
self.head2 = self.build_conv_block(128,
128,
padding=2,
dilation=2,
n_conv=1)
self.head4 = self.build_conv_block(128,
128,
padding=4,
dilation=4,
n_conv=1)
self.bbox_head, self.cls_head = self.build_det_head()
self.softmax = nn.Softmax(dim=1)
def __init__(self):
super(light_DSFD, self).__init__()
self.conv1 = CRelu(3, 32, kernel_size=7, stride=4, padding=3)
self.conv3 = CRelu(64, 64, kernel_size=5, stride=2, padding=2)
self.inception1 = Inception2d(64)
self.inception2 = Inception2d(64)
self.inception3 = Inception2d(128)
self.inception4 = Inception2d(128)
self.maxpool1 = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
self.maxpool2 = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
self.conv5_1 = BasicConv2d(128, 128, kernel_size=1, stride=1, padding=0)
self.conv5_2 = BasicConv2d(128, 256, kernel_size=3, stride=2, padding=1)
self.conv6_1 = BasicConv2d(256, 128, kernel_size=1, stride=1, padding=0)
self.conv6_2 = BasicConv2d(128, 256, kernel_size=3, stride=2, padding=1)
fpn_in = [64, 64, 128, 128, 256, 256]
cpm_in = [64, 64, 64, 64, 64, 64]
fpn_channel = 64
cpm_channels = 64
output_channels = cpm_in
# fpn
self.smooth3 = nn.Conv2d( fpn_channel, fpn_channel, kernel_size=1, stride=1, padding=0)
self.smooth2 = nn.Conv2d( fpn_channel, fpn_channel, kernel_size=1, stride=1, padding=0)
self.smooth1 = nn.Conv2d( fpn_channel, fpn_channel, kernel_size=1, stride=1, padding=0)
self.upsample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False)
self.eltmul = nn.EltMul()
self.latlayer6 = nn.Conv2d( fpn_in[5], fpn_channel, kernel_size=1, stride=1, padding=0)
self.latlayer5 = nn.Conv2d( fpn_in[4], fpn_channel, kernel_size=1, stride=1, padding=0)
self.latlayer4 = nn.Conv2d( fpn_in[3], fpn_channel, kernel_size=1, stride=1, padding=0)
self.latlayer3 = nn.Conv2d( fpn_in[2], fpn_channel, kernel_size=1, stride=1, padding=0)
self.latlayer2 = nn.Conv2d( fpn_in[1], fpn_channel, kernel_size=1, stride=1, padding=0)
self.latlayer1 = nn.Conv2d( fpn_in[0], fpn_channel, kernel_size=1, stride=1, padding=0)
# cpm
self.cpm1 = Inception2d(cpm_in[0])
self.cpm2 = Inception2d(cpm_in[1])
self.cpm3 = Inception2d(cpm_in[2])
self.cpm4 = Inception2d(cpm_in[3])
self.cpm5 = Inception2d(cpm_in[4])
self.cpm6 = Inception2d(cpm_in[5])
face_head = face_multibox(output_channels, [1, 1, 1, 1, 1, 1], 2 , cpm_channels)
self.loc = nn.ModuleList(face_head[0])
self.conf = nn.ModuleList(face_head[1])
self.softmax = nn.Softmax(dim=-1)
def __init__(self):
super(S3FD, self).__init__()
# backbone
self.vgg16 = nn.ModuleList(make_layers(vgg_cfgs['D']))
# s3fd specific
self.conv_fc6 = nn.Conv2d(512, 1024, 3, 1, 1)
self.relu_fc6 = nn.ReLU()
self.conv_fc7 = nn.Conv2d(1024, 1024, 1, 1, 0)
self.relu_fc7 = nn.ReLU()
self.conv6_1 = nn.Conv2d(1024, 256, 1, 1, 0)
self.relu_conv6_1 = nn.ReLU()
self.conv6_2 = nn.Conv2d(256, 512, 3, 2, 1)
self.relu_conv6_2 = nn.ReLU()
self.conv7_1 = nn.Conv2d(512, 128, 1, 1, 0)
self.relu_conv7_1 = nn.ReLU()
self.conv7_2 = nn.Conv2d(128, 256, 3, 2, 1)
self.relu_conv7_2 = nn.ReLU()
self.l2norm_conv3_3 = nn.L2Norm2d(256, 10)
self.l2norm_conv4_3 = nn.L2Norm2d(512, 8)
self.l2norm_conv5_3 = nn.L2Norm2d(512, 5)
# Detection Head - mbox_loc
self.mbox_loc_conv3_3_norm = nn.Conv2d(256, 4, 3, 1, 1)
self.mbox_loc_conv4_3_norm = nn.Conv2d(512, 4, 3, 1, 1)
self.mbox_loc_conv5_3_norm = nn.Conv2d(512, 4, 3, 1, 1)
self.mbox_loc_conv_fc7 = nn.Conv2d(1024, 4, 3, 1, 1)
self.mbox_loc_conv6_2 = nn.Conv2d(512, 4, 3, 1, 1)
self.mbox_loc_conv7_2 = nn.Conv2d(256, 4, 3, 1, 1)
# Detection Head - mbox_conf
self.mbox_conf_conv3_3_norm = nn.Conv2d(
256, 4, 3, 1, 1) # 4->2 through maxout at channels 0~2
self.mbox_conf_conv4_3_norm = nn.Conv2d(512, 2, 3, 1, 1)
self.mbox_conf_conv5_3_norm = nn.Conv2d(512, 2, 3, 1, 1)
self.mbox_conf_conv_fc7 = nn.Conv2d(1024, 2, 3, 1, 1)
self.mbox_conf_conv6_2 = nn.Conv2d(512, 2, 3, 1, 1)
self.mbox_conf_conv7_2 = nn.Conv2d(256, 2, 3, 1, 1)
# Detection Head - mbox_conf - softmax
self.softmax = nn.Softmax(dim=-1)
def __init__(self, mode='slim'):
super(ULFG, self).__init__()
self.mode = mode
self.base_channel = 8 * 2
self.backbone = nn.Sequential(
_conv_bn(3, self.base_channel, 2), # 160*120
_conv_dw(self.base_channel, self.base_channel * 2, 1),
_conv_dw(self.base_channel * 2, self.base_channel * 2, 2), # 80*60
_conv_dw(self.base_channel * 2, self.base_channel * 2, 1),
_conv_dw(self.base_channel * 2, self.base_channel * 4, 2), # 40*30
_conv_dw(self.base_channel * 4, self.base_channel * 4, 1),
_conv_dw(self.base_channel * 4, self.base_channel * 4, 1),
_conv_dw(self.base_channel * 4, self.base_channel * 4, 1),
_conv_dw(self.base_channel * 4, self.base_channel * 8, 2), # 20*15
_conv_dw(self.base_channel * 8, self.base_channel * 8, 1),
_conv_dw(self.base_channel * 8, self.base_channel * 8, 1),
_conv_dw(self.base_channel * 8, self.base_channel * 16, 2), # 10*8
_conv_dw(self.base_channel * 16, self.base_channel * 16, 1)
)
if self.mode == 'rfb':
self.backbone[7] = BasicRFB(self.base_channel * 4, self.base_channel * 4, stride=1, scale=1.0)
self.source_layer_indexes = [8, 11, 13]
self.extras = nn.Sequential(
nn.Conv2d(in_channels=self.base_channel * 16, out_channels=self.base_channel * 4, kernel_size=1),
nn.ReLU(),
_seperable_conv2d(in_channels=self.base_channel * 4, out_channels=self.base_channel * 16, kernel_size=3, stride=2, padding=1),
nn.ReLU()
)
self.regression_headers = nn.ModuleList([
_seperable_conv2d(in_channels=self.base_channel * 4, out_channels=3 * 4, kernel_size=3, padding=1),
_seperable_conv2d(in_channels=self.base_channel * 8, out_channels=2 * 4, kernel_size=3, padding=1),
_seperable_conv2d(in_channels=self.base_channel * 16, out_channels=2 * 4, kernel_size=3, padding=1),
nn.Conv2d(in_channels=self.base_channel * 16, out_channels=3 * 4, kernel_size=3, padding=1)
])
self.classification_headers = nn.ModuleList([
_seperable_conv2d(in_channels=self.base_channel * 4, out_channels=3 * 2, kernel_size=3, padding=1),
_seperable_conv2d(in_channels=self.base_channel * 8, out_channels=2 * 2, kernel_size=3, padding=1),
_seperable_conv2d(in_channels=self.base_channel * 16, out_channels=2 * 2, kernel_size=3, padding=1),
nn.Conv2d(in_channels=self.base_channel * 16, out_channels=3 * 2, kernel_size=3, padding=1)
])
self.softmax = nn.Softmax(dim=2)
def __init__(self):
super(YuFaceDetectNet, self).__init__()
self.num_classes = 2
self.model1 = Conv_2layers(3, 32, 16, 2)
self.model2 = Conv_2layers(16, 32, 32, 1)
self.model3 = Conv_3layers(32, 64, 32, 64, 1)
self.model4 = Conv_3layers(64, 128, 64, 128, 1)
self.model5 = Conv_3layers(128, 256, 128, 256, 1)
self.model6 = Conv_3layers(256, 256, 256, 256, 1)
self.pool1 = nn.MaxPool2d(2, 2)
self.pool2 = nn.MaxPool2d(2, 2)
self.pool3 = nn.MaxPool2d(2, 2)
self.pool4 = nn.MaxPool2d(2, 2)
self.pool5 = nn.MaxPool2d(2, 2)
self.loc, self.conf = self.multibox(self.num_classes)
self.softmax = nn.Softmax(dim=-1)
def __init__(self):
super(SFA, self).__init__()
# M0, M1, M2, M3
self.ssh_in_channels = [128, 128, 512, 512]
self.ssh_out_channels = [256, 256, 512, 512]
self.conv1 = self.build_conv_block(in_channels=3, out_channels=64, n_conv=2, with_pool=False)
self.conv2 = self.build_conv_block(in_channels=64, out_channels=128, n_conv=2, with_pool=True)
self.conv3 = self.build_conv_block(in_channels=128, out_channels=256, n_conv=3, with_pool=True)
self.conv4 = self.build_conv_block(in_channels=256, out_channels=512, n_conv=3, with_pool=True)
self.conv5 = self.build_conv_block(in_channels=512, out_channels=512, n_conv=3, with_pool=True)
# M3
self.pool6 = nn.MaxPool2d(2, 2)
self.m3 = SSH(self.ssh_in_channels[3], self.ssh_out_channels[3], index=3)
# M2
self.m2 = SSH(self.ssh_in_channels[2], self.ssh_out_channels[2], index=2)
# share by M1 and M2
self.conv4_128 = self.build_conv_block(in_channels=512, out_channels=128, kernel_size=1, padding=0, n_conv=1, with_pool=False)
# M1
self.conv5_128 = self.build_conv_block(in_channels=512, out_channels=128, kernel_size=1, padding=0, n_conv=1, with_pool=False)
self.conv5_128_up = nn.ConvTranspose2d(in_channels=128, out_channels=128, kernel_size=4, stride=2, padding=1, groups=128, bias=False)
self.conv4_fuse = nn.EltAdd()
self.conv4_fuse_final = self.build_conv_block(in_channels=128, out_channels=128, kernel_size=3, padding=1, n_conv=1, with_pool=False)
self.m1 = SSH(self.ssh_in_channels[1], self.ssh_out_channels[1], index=1)
# M0
self.conv3_128 = self.build_conv_block(in_channels=256, out_channels=128, kernel_size=1, padding=0, n_conv=1, with_pool=False)
self.conv4_128_up = nn.ConvTranspose2d(in_channels=128, out_channels=128, kernel_size=4, stride=2, padding=1, groups=128, bias=False)
self.conv3_fuse = nn.EltAdd()
self.conv3_fuse_final = self.build_conv_block(in_channels=128, out_channels=128, kernel_size=3, padding=1, n_conv=1, with_pool=False)
self.m0 = SSH(self.ssh_in_channels[0], self.ssh_out_channels[0], index=0)
# detection heads
self.bbox_head, self.cls_head = self.build_detect_head()
self.softmax = nn.Softmax(dim=-1)
def __init__(self):
super(FaceBoxes, self).__init__()
self.num_classes = 2
self.conv1 = CRelu(3, 24, kernel_size=7, stride=4, padding=3)
self.maxpool1 = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
self.conv2 = CRelu(48, 64, kernel_size=5, stride=2, padding=2)
self.maxpool2 = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
self.inception1 = Inception()
self.inception2 = Inception()
self.inception3 = Inception()
self.conv3_1 = BasicConv2d(128,
128,
kernel_size=1,
stride=1,
padding=0)
self.conv3_2 = BasicConv2d(128,
256,
kernel_size=3,
stride=2,
padding=1)
self.conv4_1 = BasicConv2d(256,
128,
kernel_size=1,
stride=1,
padding=0)
self.conv4_2 = BasicConv2d(128,
256,
kernel_size=3,
stride=2,
padding=1)
self.loc, self.conf = self.multibox(self.num_classes)
self.softmax = nn.Softmax(dim=-1)
def __init__(self, in_channels):
super(DetBlock, self).__init__()
self.in_channels = in_channels
self.det_channels = 128
self.det_conv = nn.Conv2d(in_channels=self.in_channels,
out_channels=self.det_channels,
kernel_size=1,
stride=1,
padding=0)
self.det_relu = nn.ReLU()
self.bbox_conv = nn.Conv2d(in_channels=self.det_channels,
out_channels=self.det_channels,
kernel_size=1,
stride=1,
padding=0)
self.bbox_relu = nn.ReLU()
self.bbox_out_conv = nn.Conv2d(in_channels=self.det_channels,
out_channels=4,
kernel_size=1,
stride=1,
padding=0)
self.score_conv = nn.Conv2d(in_channels=self.det_channels,
out_channels=self.det_channels,
kernel_size=1,
stride=1,
padding=0)
self.score_relu = nn.ReLU()
self.score_out_conv = nn.Conv2d(in_channels=self.det_channels,
out_channels=2,
kernel_size=1,
stride=1,
padding=0)
self.softmax = nn.Softmax(dim=1)
def __init__(self):
super(DSFD, self).__init__()
self.size = 640
self.num_classes = 2
######
# build backbone
######
resnet152 = vision.models.resnet152()
self.layer1 = nn.Sequential(resnet152.conv1, resnet152.bn1,
resnet152.relu, resnet152.maxpool,
resnet152.layer1)
self.layer2 = nn.Sequential(resnet152.layer2)
self.layer3 = nn.Sequential(resnet152.layer3)
self.layer4 = nn.Sequential(resnet152.layer4)
self.layer5 = nn.Sequential(*[
nn.Conv2d(2048, 512, kernel_size=1),
nn.BatchNorm2d(512),
nn.ReLU(inplace=True),
nn.Conv2d(512, 512, kernel_size=3, padding=1, stride=2),
nn.BatchNorm2d(512),
nn.ReLU(inplace=True)
])
self.layer6 = nn.Sequential(*[
nn.Conv2d(
512,
128,
kernel_size=1,
),
nn.BatchNorm2d(128),
nn.ReLU(inplace=True),
nn.Conv2d(128, 256, kernel_size=3, padding=1, stride=2),
nn.BatchNorm2d(256),
nn.ReLU(inplace=True)
])
######
# dsfd specific layers
######
output_channels = [256, 512, 1024, 2048, 512, 256]
# fpn
fpn_in = output_channels
self.latlayer3 = nn.Conv2d(fpn_in[3],
fpn_in[2],
kernel_size=1,
stride=1,
padding=0)
self.latlayer2 = nn.Conv2d(fpn_in[2],
fpn_in[1],
kernel_size=1,
stride=1,
padding=0)
self.latlayer1 = nn.Conv2d(fpn_in[1],
fpn_in[0],
kernel_size=1,
stride=1,
padding=0)
self.smooth3 = nn.Conv2d(fpn_in[2],
fpn_in[2],
kernel_size=1,
stride=1,
padding=0)
self.smooth2 = nn.Conv2d(fpn_in[1],
fpn_in[1],
kernel_size=1,
stride=1,
padding=0)
self.smooth1 = nn.Conv2d(fpn_in[0],
fpn_in[0],
kernel_size=1,
stride=1,
padding=0)
self.upsample = nn.Upsample(scale_factor=2,
mode='bilinear',
align_corners=False)
self.eltmul = nn.EltMul()
# fem
cpm_in = output_channels
self.cpm3_3 = FEM(cpm_in[0])
self.cpm4_3 = FEM(cpm_in[1])
self.cpm5_3 = FEM(cpm_in[2])
self.cpm7 = FEM(cpm_in[3])
self.cpm6_2 = FEM(cpm_in[4])
self.cpm7_2 = FEM(cpm_in[5])
# pa
cfg_mbox = [1, 1, 1, 1, 1, 1]
head = pa_multibox(output_channels, cfg_mbox, self.num_classes)
# detection head
self.loc = nn.ModuleList(head[0])
self.conf = nn.ModuleList(head[1])
self.softmax = nn.Softmax(dim=-1)
}
test_on(relu)
######
# test on Sigmoid
######
sigmoid = {
'layers': [
nn.Sigmoid() # same shape
],
'ins': [TensorSize([1, 1, 56, 56])],
'out_shape': [TensorSize([1, 1, 56, 56])],
'out_flops': [9408]
}
test_on(sigmoid)
######
# test on Softmax
######
softmax = {
'layers': [
nn.Softmax(dim=-1) # same shape
],
'ins': [TensorSize([1, 4185, 2])],
'out_shape': [TensorSize([1, 4185, 2])],
'out_flops': [25108]
}
test_on(softmax)