def get_flop(self, kernel_size, stride, in_channel, out_channel, affine, in_h, in_w): cal_flop = flops.count_ConvBNReLU_flop(in_h, in_w, in_channel, self.hidden_dim, [1, 1], False, is_affine=affine) cal_flop += flops.count_ConvBNReLU_flop(in_h, in_w, self.hidden_dim, self.hidden_dim, kernel_size, False, stride=stride, is_affine=affine) in_h, in_w = in_h // stride, in_w // stride cal_flop += flops.count_Conv_flop(in_h, in_w, self.hidden_dim, out_channel, kernel_size=[1, 1], is_bias=False) cal_flop += flops.count_BN_flop(in_h, in_w, out_channel, affine) return cal_flop
def get_flop(self, kernel_size, stride, in_channel, out_channel, affine, in_h, in_w): cal_flop = flops.count_Conv_flop(in_h, in_w, in_channel, in_channel, kernel_size, False, stride, groups=in_channel) in_h, in_w = in_h // stride, in_w // stride cal_flop += flops.count_ConvBNReLU_flop(in_h, in_w, in_channel, in_channel, kernel_size=[1, 1], is_bias=False, is_affine=affine) # stack 2 separate depthwise-conv. cal_flop += flops.count_Conv_flop(in_h, in_w, in_channel, in_channel, kernel_size, False, stride=1, groups=in_channel) cal_flop += flops.count_ConvBNReLU_flop(in_h, in_w, in_channel, in_channel, kernel_size=[1, 1], is_bias=False, is_affine=affine) # stack 3 separate depthwise-conv. cal_flop += flops.count_Conv_flop(in_h, in_w, in_channel, in_channel, kernel_size, False, stride=1, groups=in_channel) cal_flop += flops.count_Conv_flop(in_h, in_w, in_channel, out_channel, kernel_size=[1, 1], is_bias=False) cal_flop += flops.count_BN_flop(in_h, in_w, out_channel, affine) return cal_flop
def __init__(self, cfg, input_shape: Dict[str, ShapeSpec]): super().__init__() # fmt: off self.in_features = cfg.MODEL.SEM_SEG_HEAD.IN_FEATURES feature_strides = {k: v.stride for k, v in input_shape.items()} # noqa:F841 feature_channels = {k: v.channels for k, v in input_shape.items()} feature_resolution = { k: np.array([v.height, v.width]) for k, v in input_shape.items() } self.ignore_value = cfg.MODEL.SEM_SEG_HEAD.IGNORE_VALUE num_classes = cfg.MODEL.SEM_SEG_HEAD.NUM_CLASSES norm = cfg.MODEL.SEM_SEG_HEAD.NORM self.loss_weight = cfg.MODEL.SEM_SEG_HEAD.LOSS_WEIGHT self.cal_flops = cfg.MODEL.CAL_FLOPS self.real_flops = 0.0 # fmt: on self.layer_decoder_list = nn.ModuleList() # set affine in BatchNorm if 'Sync' in norm: affine = True else: affine = False # use simple decoder for _feat in self.in_features: res_size = feature_resolution[_feat] in_channel = feature_channels[_feat] if _feat == 'layer_0': out_channel = in_channel else: out_channel = in_channel // 2 conv_1x1 = Conv2d(in_channel, out_channel, kernel_size=1, stride=1, padding=0, bias=False, norm=get_norm(norm, out_channel), activation=nn.ReLU()) self.real_flops += cal_op_flops.count_ConvBNReLU_flop( res_size[0], res_size[1], in_channel, out_channel, [1, 1], is_affine=affine) self.layer_decoder_list.append(conv_1x1) # using Kaiming init for layer in self.layer_decoder_list: weight_init.kaiming_init_module(layer, mode='fan_in') in_channel = feature_channels['layer_0'] # the output layer self.predictor = Conv2d(in_channels=in_channel, out_channels=num_classes, kernel_size=3, stride=1, padding=1) self.real_flops += cal_op_flops.count_Conv_flop( feature_resolution['layer_0'][0], feature_resolution['layer_0'][1], in_channel, num_classes, [3, 3]) # using Kaiming init weight_init.kaiming_init_module(self.predictor, mode='fan_in')
def __init__(self, in_channels=3, mid_channels=64, out_channels=64, input_res=None, sept_stem=True, norm="BN", affine=True): """ Build basic STEM for Dynamic Network. Args: norm (str or callable): a callable that takes the number of channels and return a `nn.Module`, or a pre-defined string (one of {"FrozenBN", "BN", "GN"}). """ super().__init__() self.real_flops = 0.0 # start with 3 stem layers down-sampling by 4. self.stem_1 = Conv2d(in_channels, mid_channels, kernel_size=3, stride=2, bias=False, norm=get_norm(norm, mid_channels), activation=nn.ReLU()) self.real_flops += cal_op_flops.count_ConvBNReLU_flop(input_res[0], input_res[1], 3, mid_channels, [3, 3], stride=2, is_affine=affine) # stem 2 input_res = input_res // 2 if not sept_stem: self.stem_2 = Conv2d(mid_channels, mid_channels, kernel_size=3, stride=1, padding=1, bias=False, norm=get_norm(norm, mid_channels), activation=nn.ReLU()) self.real_flops += cal_op_flops.count_ConvBNReLU_flop( input_res[0], input_res[1], mid_channels, mid_channels, [3, 3], is_affine=affine) else: self.stem_2 = nn.Sequential( Conv2d(mid_channels, mid_channels, kernel_size=3, stride=1, padding=1, groups=mid_channels, bias=False), Conv2d(mid_channels, mid_channels, kernel_size=1, stride=1, padding=0, bias=False, norm=get_norm(norm, mid_channels), activation=nn.ReLU())) self.real_flops += ( cal_op_flops.count_Conv_flop(input_res[0], input_res[1], mid_channels, mid_channels, [3, 3], groups=mid_channels) + cal_op_flops.count_ConvBNReLU_flop(input_res[0], input_res[1], mid_channels, mid_channels, [1, 1], is_affine=affine)) # stem 3 if not sept_stem: self.stem_3 = Conv2d(mid_channels, out_channels, kernel_size=3, stride=2, padding=1, bias=False, norm=get_norm(norm, out_channels), activation=nn.ReLU()) self.real_flops += cal_op_flops.count_ConvBNReLU_flop( input_res[0], input_res[1], mid_channels, out_channels, [3, 3], stride=2, is_affine=affine) else: self.stem_3 = nn.Sequential( Conv2d(mid_channels, mid_channels, kernel_size=3, stride=2, padding=1, groups=mid_channels, bias=False), Conv2d(mid_channels, out_channels, kernel_size=1, padding=0, bias=False, norm=get_norm(norm, out_channels), activation=nn.ReLU())) self.real_flops += ( cal_op_flops.count_Conv_flop(input_res[0], input_res[1], mid_channels, mid_channels, [3, 3], stride=2, groups=mid_channels) + cal_op_flops.count_ConvBNReLU_flop(input_res[0] // 2, input_res[1] // 2, mid_channels, out_channels, [1, 1], is_affine=affine)) self.out_res = input_res // 2 self.out_cha = out_channels # using Kaiming init for layer in [self.stem_1, self.stem_2, self.stem_3]: weight_init.kaiming_init_module(layer, mode='fan_in')
def __init__(self, C_in, C_out, norm, allow_up, allow_down, input_size, cell_type, cal_flops=True, using_gate=False, small_gate=False, gate_bias=1.5, affine=True): super(Cell, self).__init__() self.channel_in = C_in self.channel_out = C_out self.allow_up = allow_up self.allow_down = allow_down self.cal_flops = cal_flops self.using_gate = using_gate self.small_gate = small_gate self.cell_ops = Mixed_OP(inplanes=self.channel_in, outplanes=self.channel_out, stride=1, cell_type=cell_type, norm=norm, affine=affine, input_size=input_size) self.cell_flops = self.cell_ops.flops # resolution keep self.res_keep = nn.ReLU() self.res_keep_flops = cal_op_flops.count_ReLU_flop( input_size[0], input_size[1], self.channel_out) # resolution up and dim down if self.allow_up: self.res_up = nn.Sequential( nn.ReLU(), Conv2d(self.channel_out, self.channel_out // 2, kernel_size=1, stride=1, padding=0, bias=False, norm=get_norm(norm, self.channel_out // 2), activation=nn.ReLU())) # calculate Flops self.res_up_flops = cal_op_flops.count_ReLU_flop( input_size[0], input_size[1], self.channel_out) + cal_op_flops.count_ConvBNReLU_flop( input_size[0], input_size[1], self.channel_out, self.channel_out // 2, [1, 1], is_affine=affine) # using Kaiming init weight_init.kaiming_init_module(self.res_up, mode='fan_in') # resolution down and dim up if self.allow_down: self.res_down = nn.Sequential( nn.ReLU(), Conv2d(self.channel_out, 2 * self.channel_out, kernel_size=1, stride=2, padding=0, bias=False, norm=get_norm(norm, 2 * self.channel_out), activation=nn.ReLU())) # calculate Flops self.res_down_flops = cal_op_flops.count_ReLU_flop( input_size[0], input_size[1], self.channel_out) + cal_op_flops.count_ConvBNReLU_flop( input_size[0], input_size[1], self.channel_out, 2 * self.channel_out, [1, 1], stride=2, is_affine=affine) # using Kaiming init weight_init.kaiming_init_module(self.res_down, mode='fan_in') if self.allow_up and self.allow_down: self.gate_num = 3 elif self.allow_up or self.allow_down: self.gate_num = 2 else: self.gate_num = 1 if self.using_gate: self.gate_conv_beta = nn.Sequential( Conv2d(self.channel_in, self.channel_in // 2, kernel_size=1, stride=1, padding=0, bias=False, norm=get_norm(norm, self.channel_in // 2), activation=nn.ReLU()), nn.AdaptiveAvgPool2d((1, 1)), Conv2d(self.channel_in // 2, self.gate_num, kernel_size=1, stride=1, padding=0, bias=True)) if self.small_gate: input_size = input_size // 4 self.gate_flops = cal_op_flops.count_ConvBNReLU_flop( input_size[0], input_size[1], self.channel_in, self.channel_in // 2, [1, 1], is_affine=affine) + cal_op_flops.count_Pool2d_flop( input_size[0], input_size[1], self.channel_in // 2, [1, 1], 1) + cal_op_flops.count_Conv_flop( 1, 1, self.channel_in // 2, self.gate_num, [1, 1]) # using Kaiming init and predefined bias for gate weight_init.kaiming_init_module(self.gate_conv_beta, mode='fan_in', bias=gate_bias) else: self.register_buffer('gate_weights_beta', torch.ones(1, self.gate_num, 1, 1).cuda()) self.gate_flops = 0.0