def forward(self, z, is_training=True, test_local_stats=True): batch_norm_args = {'track_running_stats': test_local_stats} method = self._method # Cycle over the encoder shapes backwards, to build a symmetrical decoder. enc_conv_shapes = self._enc_conv_shapes[::-1] strides = self._dec_up_strides # We store the heights and widths of the encoder feature maps that are # unique, i.e., the ones right after a layer with stride != 1. These will be # used as a target to potentially crop the upsampled feature maps. unique_hw = np.unique([(el[1], el[2]) for el in enc_conv_shapes], axis=0) unique_hw = unique_hw.tolist()[::-1] unique_hw.pop() # Drop the initial shape layers = [z] upsample_mlp_flat = self.dec_mlp(z) if self._use_bn: upsample_mlp_flat = nn.BatchNorm1d( upsample_mlp_flat.shape[1], **batch_norm_args)(upsample_mlp_flat) layers.append(upsample_mlp_flat) upsample = upsample_mlp_flat.view(enc_conv_shapes[0]) layers.append(upsample) for i, (conv_layer, stride_i) in enumerate(zip(self._conv_layers, strides), 1): if method != 'deconv' and stride_i > 1: upsample = Image.resize( upsample, [stride_i * el for el in upsample.shape[1:3]], method=method) padding = get_same_padding(self._kernel_size, upsample.shape, stride_i) upsample_padded = F.pad(upsample, padding, "constant", 0) upsample = conv_layer(upsample_padded) upsample = self._activation(upsample) if self._use_bn: upsample = nn.BatchNorm2d(upsample.shape[1], **batch_norm_args)(upsample) if stride_i > 1: hw = unique_hw.pop() upsample = crop_layer(upsample, hw) layers.append(upsample) # Final layer, no upsampling. padding = get_same_padding(self._kernel_size, upsample.shape, 1) upsample_padded = F.pad(upsample, padding, "constant", 0) x_logits = self.logits_layer(upsample_padded) if self._use_bn: x_logits = nn.BatchNorm2d(x_logits.shape[1], **batch_norm_args)(x_logits) layers.append(x_logits) logging.info('%s upsampling module layer shapes', self._method_str) logging.info('\n'.join([str(v.shape) for v in layers])) return x_logits
def make_model(pool_op, shape, pool_size, strides, padding, dtype, scale, zero_point, relu_type): """Return a model and any parameters it may have""" op = relay.var("input", shape=shape, dtype=dtype) pad_ = (0, 0, 0, 0) if padding == "SAME": dilation = (1, 1) pad_ = get_same_padding((shape[1], shape[2]), pool_size, dilation, strides) op = relay.nn.pad( op, pad_width=[(0, 0), (pad_[0], pad_[2]), (pad_[1], pad_[3]), (0, 0)], pad_value=zero_point, pad_mode="constant", ) if pool_op == relay.nn.avg_pool2d: op = relay.cast(op, "int32") op = pool_op(op, pool_size=pool_size, strides=strides, padding=pad_, ceil_mode=True, layout="NHWC") if pool_op == relay.nn.avg_pool2d: op = relay.cast(op, dtype) op = make_qnn_relu(op, relu_type, scale, zero_point, dtype) return op
def weight_op(self): padding = get_same_padding(self.kernel_size) if isinstance(padding, int): padding *= self.dilation else: padding[0] *= self.dilation padding[1] *= self.dilation weight_dict = OrderedDict() weight_dict['depth_conv'] = nn.Conv2d(self.in_channels, self.in_channels, kernel_size=self.kernel_size, stride=self.stride, padding=padding, dilation=self.dilation, groups=self.in_channels, bias=False) weight_dict['point_conv'] = nn.Conv2d(self.in_channels, self.out_channels, kernel_size=1, groups=self.groups, bias=self.bias) if self.has_shuffle and self.groups > 1: weight_dict['shuffle'] = ShuffleLayer(self.groups) return weight_dict
def weight_op(self): padding = get_same_padding(self.kernel_size) if isinstance(padding, int): padding *= self.dilation else: padding[0] *= self.dilation padding[1] *= self.dilation weight_dict = OrderedDict() # weight_dict['conv'] = nn.Conv2d( # self.in_channels, self.out_channels, kernel_size=self.kernel_size, stride=self.stride, padding=padding, # dilation=self.dilation, groups=self.groups, bias=self.bias # ) weight_dict['conv'] = LsqConv(self.in_channels, self.out_channels, kernel_size=self.kernel_size, quan_name_w='lsq', quan_name_a='lsq', nbit_w=self.nbit_w, nbit_a=self.nbit_a, stride=self.stride, padding=padding, dilation=self.dilation, groups=self.groups, bias=self.bias) if self.has_shuffle and self.groups > 1: weight_dict['shuffle'] = ShuffleLayer(self.groups) return weight_dict
def __init__(self, in_channels, out_channels, pool_type, kernel_size=2, stride=2, use_bn=False, act_func=None, dropout_rate=0, ops_order='weight_bn_act'): super(PoolingLayer, self).__init__(in_channels, out_channels, use_bn, act_func, dropout_rate, ops_order) self.pool_type = pool_type self.kernel_size = kernel_size self.stride = stride if self.stride == 1: # same padding if `stride == 1` padding = get_same_padding(self.kernel_size) else: padding = 0 if self.pool_type == 'avg': self.pool = nn.AvgPool2d(self.kernel_size, stride=self.stride, padding=padding, count_include_pad=False) elif self.pool_type == 'max': self.pool = nn.MaxPool2d(self.kernel_size, stride=self.stride, padding=padding) else: raise NotImplementedError
def forward(self, X): self.N, C, X_n, _ = X.shape self.X_shape = X.shape if self.same_padding: self.padding = utils.get_same_padding(X_n, self.kernel_size, self.stride) else: self.padding = 0 _, self.new_X_n, Xcol = standfordutils.im2col(X, self.kernel_size, self.kernel_size, self.padding, self.stride) self.S = gates.ConvMultGate() self.Z = gates.ConvAddGate() S_out = self.S.forward(self.W, Xcol) Z_out = self.Z.forward(S_out, self.B) Z_out = Z_out.reshape(self.N, self.n_kernels, self.new_X_n, self.new_X_n) if self.logging: print("[+] CONVOLUTION_OUT:", Z_out.shape) return Z_out
def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, expand_ratio=6, *args, **kwargs): super(MBInvertedConvLayer, self).__init__() self.in_channels = in_channels self.out_channels = out_channels self.kernel_size = kernel_size self.stride = stride self.expand_ratio = expand_ratio if self.expand_ratio > 1: feature_dim = round(in_channels * self.expand_ratio) self.inverted_bottleneck = nn.Sequential( OrderedDict({ 'conv': nn.Conv2d(in_channels, feature_dim, 1, 1, 0, bias=False), 'bn': nn.BatchNorm2d(feature_dim), 'relu': nn.ReLU6(inplace=True), })) else: feature_dim = in_channels self.inverted_bottleneck = None # depthwise convolution pad = get_same_padding(self.kernel_size) self.depth_conv = nn.Sequential( OrderedDict({ 'conv': nn.Conv2d(feature_dim, feature_dim, kernel_size, stride, pad, groups=feature_dim, bias=False), 'bn': nn.BatchNorm2d(feature_dim), 'relu': nn.ReLU6(inplace=True), })) # pointwise linear self.point_linear = OrderedDict({ 'conv': nn.Conv2d(feature_dim, out_channels, 1, 1, 0, bias=False), 'bn': nn.BatchNorm2d(out_channels), }) self.point_linear = nn.Sequential(self.point_linear)
def export(self) -> nn.Module: module = getattr(nn, "Conv{}d".format(self._ndim))( self.active_in_channels, self.active_in_channels, self.kernel_size[0], stride=self.stride, padding=get_same_padding(self.kernel_size[0]) * self.dilation[0], dilation=self.dilation, groups=self.active_in_channels, bias=self.bias is not None, padding_mode=self.padding_mode, ) module.load_state_dict(self.active_state_dict()) return module
def forward(self, x: torch.Tensor) -> torch.Tensor: self.active_in_channels = x.shape[1] if self.padding_mode != "zeros": raise NotImplementedError else: active_weight = self.active_weight return getattr(F, "conv{}d".format(self._ndim))( x, active_weight, self.active_bias, stride=self.stride, padding=get_same_padding(int(active_weight.size(2))) * self.dilation[0], dilation=self.dilation, groups=self.active_in_channels, )
def weight_op(self): if self.stride == 1: # same padding if `stride == 1` padding = get_same_padding(self.kernel_size) else: padding = 0 weight_dict = OrderedDict() if self.pool_type == 'avg': weight_dict['pool'] = nn.AvgPool2d( self.kernel_size, stride=self.stride, padding=padding, count_include_pad=False ) elif self.pool_type == 'max': weight_dict['pool'] = nn.MaxPool2d(self.kernel_size, stride=self.stride, padding=padding) else: raise NotImplementedError return weight_dict
def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, dilation=1, groups=1, bias=False, has_shuffle=False, use_bn=True, act_func='relu', dropout_rate=0, ops_order='weight_bn_act', *args, **kwargs): super(DepthConvLayer, self).__init__(in_channels, out_channels, use_bn, act_func, dropout_rate, ops_order) self.kernel_size = kernel_size self.stride = stride self.dilation = dilation self.groups = groups self.bias = bias self.has_shuffle = has_shuffle padding = get_same_padding(self.kernel_size) if isinstance(padding, int): padding *= self.dilation else: padding[0] *= self.dilation padding[1] *= self.dilation # `kernel_size`, `stride`, `padding`, `dilation` can either be `int` or `tuple` of int self.depth_conv = nn.Conv2d(in_channels, in_channels, kernel_size=self.kernel_size, stride=self.stride, padding=padding, dilation=self.dilation, groups=in_channels, bias=False) self.point_conv = nn.Conv2d(in_channels, out_channels, kernel_size=1, groups=self.groups, bias=self.bias)
def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, expand_ratio=6, mid_channels=None, act_func='relu6', use_se=False): super(MBInvertedConvLayer, self).__init__() self.in_channels = in_channels self.out_channels = out_channels self.kernel_size = kernel_size self.stride = stride self.expand_ratio = expand_ratio self.mid_channels = mid_channels self.act_func = act_func self.use_se = use_se if self.mid_channels is None: feature_dim = round(self.in_channels * self.expand_ratio) else: feature_dim = self.mid_channels if self.expand_ratio == 1: self.inverted_bottleneck = None else: self.inverted_bottleneck = nn.Sequential(OrderedDict([ ('conv', nn.Conv2d(self.in_channels, feature_dim, 1, 1, 0, bias=False)), ('bn', nn.BatchNorm2d(feature_dim)), ('act', build_activation(self.act_func, inplace=True)), ])) pad = get_same_padding(self.kernel_size) depth_conv_modules = [ ('conv', nn.Conv2d(feature_dim, feature_dim, kernel_size, stride, pad, groups=feature_dim, bias=False)), ('bn', nn.BatchNorm2d(feature_dim)), ('act', build_activation(self.act_func, inplace=True)) ] if self.use_se: depth_conv_modules.append(('se', SEModule(feature_dim))) self.depth_conv = nn.Sequential(OrderedDict(depth_conv_modules)) self.point_linear = nn.Sequential(OrderedDict([ ('conv', nn.Conv2d(feature_dim, out_channels, 1, 1, 0, bias=False)), ('bn', nn.BatchNorm2d(out_channels)), ]))
def make_model( pool_op, shape=(1, 28, 28, 12), pool_size=(3, 3), strides=(2, 2), padding="VALID", dtype="int8", scale=1, zero_point=-33, relu_type="RELU", layout="NHWC", ): """Return a model and any parameters it may have, all parameters are defaulted to known good values""" op = relay.var("input", shape=shape, dtype=dtype) pad_ = (0, 0, 0, 0) if padding == "SAME": dilation = (1, 1) pad_ = get_same_padding((shape[1], shape[2]), pool_size, dilation, strides) op = relay.nn.pad( op, pad_width=[(0, 0), (pad_[0], pad_[2]), (pad_[1], pad_[3]), (0, 0)], pad_value=zero_point, pad_mode="constant", ) if pool_op == relay.nn.avg_pool2d: op = relay.cast(op, "int32") op = pool_op(op, pool_size=pool_size, strides=strides, padding=pad_, ceil_mode=True, layout=layout) if pool_op == relay.nn.avg_pool2d: op = relay.cast(op, dtype) op = make_qnn_relu(op, relu_type, scale, zero_point, dtype) return op
def make_model( shape, kernel_shape, input_zero_point, input_scale, kernel_zero_point, kernel_scale, output_zero_point, output_scale, padding, strides, dilation, groups, dtype, kernel_dtype, out_channels, weight_format, enable_bias, relu_type, ): """Return a model and any parameters it may have""" h_index = weight_format.index("H") w_index = weight_format.index("W") kernel_h = kernel_shape[h_index] kernel_w = kernel_shape[w_index] a = relay.var("input", shape=shape, dtype=dtype) p = (0, 0, 0, 0) if padding == "SAME": p = get_same_padding((shape[1], shape[2]), (kernel_h, kernel_w), dilation, strides) a = relay.nn.pad( a, pad_width=[(0, 0), (p[0], p[2]), (p[1], p[3]), (0, 0)], pad_value=input_zero_point, pad_mode="constant", ) shape = (shape[0], shape[1] + p[0] + p[2], shape[2] + p[1] + p[3], shape[3]) weight_shape = (kernel_h, kernel_w, shape[3] // groups, out_channels) rng = np.random.default_rng(12321) w = tvm.nd.array( rng.integers( np.iinfo(kernel_dtype).min, high=np.iinfo(kernel_dtype).max, size=weight_shape, dtype=kernel_dtype, )) weight_const = relay.const(w, kernel_dtype) conv = relay.qnn.op.conv2d( a, weight_const, input_zero_point=relay.const(input_zero_point, "int32"), kernel_zero_point=relay.const(kernel_zero_point, "int32"), input_scale=relay.const(input_scale, "float32"), kernel_scale=relay.const(kernel_scale, "float32"), kernel_size=(kernel_h, kernel_w), data_layout="NHWC", kernel_layout=weight_format, dilation=dilation, strides=strides, groups=groups, channels=out_channels, padding=p, out_dtype="int32", ) b = tvm.nd.array( rng.integers(0, high=10, size=(out_channels, ), dtype="int32")) bias_const = relay.const(b, "int32") last_op = relay.nn.bias_add(conv, bias_const, axis=3) if enable_bias else conv requant_input_sc = [sc * input_scale for sc in kernel_scale] last_op = relay.qnn.op.requantize( last_op, relay.const(requant_input_sc, "float32"), relay.const(0, "int32"), relay.const(output_scale, "float32"), relay.const(output_zero_point, "int32"), out_dtype=dtype, ) last_op = make_qnn_relu(last_op, relu_type, output_scale, output_zero_point, dtype) params = {"w": w, "b": b} return last_op, params