|
- import time
- import numpy as np
- import mindspore as ms
- import mindspore.nn as nn
- import mindspore.ops.operations as P
-
- ms.set_seed(1)
- ms.set_context(mode=0)
-
-
- class GroupConv(nn.Cell):
- """
- group convolution operation.
-
- Args:
- in_channels (int): Input channels of feature map.
- out_channels (int): Output channels of feature map.
- kernel_size (int): Size of convolution kernel.
- stride (int): Stride size for the group convolution layer.
-
- Returns:
- tensor, output tensor.
- """
-
- def __init__(self, in_channels, out_channels, kernel_size, stride, pad_mode="pad", pad=0, groups=1, has_bias=False):
- super(GroupConv, self).__init__()
- assert in_channels % groups == 0 and out_channels % groups == 0
- self.groups = groups
- self.convs = nn.CellList()
- self.op_split = P.Split(axis=1, output_num=self.groups)
- self.op_concat = P.Concat(axis=1)
- self.cast = P.Cast()
- for _ in range(groups):
- self.convs.append(nn.Conv2d(in_channels // groups, out_channels // groups,
- kernel_size=kernel_size, stride=stride, has_bias=has_bias,
- padding=pad, pad_mode=pad_mode, group=1, weight_init='xavier_uniform'))
-
- def construct(self, x):
- features = self.op_split(x)
- outputs = ()
- for i in range(self.groups):
- outputs = outputs + (self.convs[i](self.cast(features[i], ms.float32)),)
- out = self.op_concat(outputs)
- return out
-
-
- for i in range(10):
-
- print("iter%d" % (i+1))
-
- start1 = time.time()
- conv1 = GroupConv(in_channels=3, out_channels=24, kernel_size=1, stride=1, pad_mode="pad", pad=0, groups=3)
- input1 = ms.Tensor(np.ones([1, 3, 224, 224]), dtype=ms.float32)
- output1 = conv1(input1)
- end1 = time.time()
- print("1.5x GroupConv time_cost:", end1 - start1)
-
- start2 = time.time()
- conv2 = nn.Conv2d(in_channels=3, out_channels=24, kernel_size=1, stride=1, pad_mode="pad", group=3)
- input2 = ms.Tensor(np.ones([1, 3, 224, 224]), dtype=ms.float32)
- output2 = conv2(input2)
- end2 = time.time()
- print("1.5x Conv2d time_cost:", end2 - start2)
-
- start3 = time.time()
- conv3 = GroupConv(in_channels=3, out_channels=48, kernel_size=1, stride=1, pad_mode="pad", pad=0, groups=3)
- input3 = ms.Tensor(np.ones([1, 3, 224, 224]), dtype=ms.float32)
- output3 = conv3(input3)
- end3 = time.time()
- print("2.0x GroupConv time_cost:", end3 - start3)
-
- start4 = time.time()
- conv4 = nn.Conv2d(in_channels=3, out_channels=48, kernel_size=1, stride=1, pad_mode="pad", group=3)
- input4 = ms.Tensor(np.ones([1, 3, 224, 224]), dtype=ms.float32)
- output4 = conv4(input4)
- end4 = time.time()
- print("2.0x Conv2d time_cost:", end4 - start4)
|