def Block(dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0., drop_path=0): mlp_hidden_dim = int(dim * mlp_ratio) return Sequential( ShortCut(Identity(), Sequential( LayerNorm(), Attention(dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop), DropPath()), mode='add'), ShortCut(Identity(), Sequential( LayerNorm(), Mlp(hidden_features=mlp_hidden_dim, out_features=dim, drop=drop), DropPath()), mode='add'))
def __init__(self, *args, axis=1, activation='relu'): """ Parameters ---------- layer_defs : object """ super(ConcateBlock, self).__init__() self.activation = get_activation(activation) self.axis = axis self.has_identity = False for i in range(len(args)): arg = args[i] if isinstance(arg, (Layer, list, dict)): if isinstance(arg, list): arg = Sequential(*arg) elif isinstance(arg, dict) and len(args) == 1: for k, v in arg.items(): if isinstance(v, Identity): self.has_identity = True self.add_module('Identity', v) else: self.add_module(k, v) elif isinstance(arg, dict) and len(args) > 1: raise ValueError( 'more than one dict argument is not support.') elif isinstance(arg, Identity): self.has_identity = True self.add_module('Identity', arg) else: self.add_module('branch{0}'.format(i + 1), arg) if len(self._modules) == 1 and self.has_identity == False: self.add_module('Identity', Identity()) self.to(self.device)
def basic_block(num_filters=64,base_width=64,strides=1,expansion = 4,conv_shortcut=False,use_bias=False,name=''): shortcut = Identity() if strides>1 or conv_shortcut is True: shortcut =Conv2d_Block((1,1),num_filters=num_filters,strides=strides,auto_pad=True,padding_mode='zero',normalization='batch',activation=None,use_bias=use_bias) return ShortCut2d(Sequential(Conv2d_Block((3,3),num_filters=num_filters,strides=strides,auto_pad=True,padding_mode='zero',normalization='batch',activation='relu',use_bias=use_bias), Conv2d_Block((3,3),num_filters=num_filters,strides=1,auto_pad=True,padding_mode='zero',normalization='batch',activation=None,use_bias=use_bias)), shortcut,activation='relu')
def se_bottleneck(num_filters=64, strides=1, expansion=4, conv_shortcut=True, use_bias=False, name=''): #width = int(num_filters * (base_width / 64.)) * 1#groups' shortcut = Identity() shortcut_name = 'Identity' if strides > 1 or conv_shortcut is True: shortcut = Conv2d_Block((1, 1), num_filters=num_filters * expansion, strides=strides, auto_pad=True, padding_mode='zero', normalization='batch', activation=None, use_bias=use_bias, name='downsample') shortcut_name = 'downsample' return ShortCut2d( { 'branch1': Sequential( Conv2d_Block((1, 1), num_filters=num_filters, strides=strides, auto_pad=True, padding_mode='zero', normalization='batch', activation='relu', use_bias=use_bias, name='conv1'), Conv2d_Block((3, 3), num_filters=num_filters, strides=1, auto_pad=True, padding_mode='zero', normalization='batch', activation='relu', use_bias=use_bias, name='conv2'), Conv2d_Block((1, 1), num_filters=num_filters * expansion, strides=1, auto_pad=True, padding_mode='zero', normalization='batch', activation=None, use_bias=use_bias, name='conv3'), SqueezeExcite(se_filters=num_filters // expansion, num_filters=num_filters * expansion, use_bias=True)), shortcut_name: shortcut }, activation='relu')
def separable_resnet_block(num_filters=64, strides=1, activation='leaky_relu', normalization='batch', use_spectral=False, dilation=1, name=''): kernal = 1 if strides == 1 else 3 return [ShortCut2d(Sequential( SeparableConv2d_Block((3, 3), depth_multiplier=1, strides=1, auto_pad=True, padding_mode='replicate', use_spectral=use_spectral, normalization=normalization, activation=activation, use_bias=False, dilation=dilation, name=name + '_0_conv'), Conv2d_Block((1, 1), depth_multiplier=2, strides=1, auto_pad=True, padding_mode='replicate', use_spectral=use_spectral, normalization=normalization, activation=activation, use_bias=False, name=name + '_1_conv')), Identity(), activation=activation, name=name), Conv2d_Block((kernal, kernal), num_filters=num_filters, strides=strides, auto_pad=True, padding_mode='replicate', use_spectral=use_spectral, normalization=normalization, activation=activation, use_bias=False, name=name + '_conv')]
def resblock_body(num_filters, num_blocks, all_narrow=True, keep_output=False, name=''): return Sequential( DarknetConv2D_BN_Mish((3, 3), num_filters, strides=(2, 2), auto_pad=False, padding=((1, 0), (1, 0)), name=name + '_preconv1'), ShortCut2d( { 1: DarknetConv2D_BN_Mish( (1, 1), num_filters // 2 if all_narrow else num_filters, name=name + '_shortconv'), 0: Sequential( DarknetConv2D_BN_Mish( (1, 1), num_filters // 2 if all_narrow else num_filters, name=name + '_mainconv'), For( range(num_blocks), lambda i: ShortCut2d( Identity(), Sequential( DarknetConv2D_BN_Mish( (1, 1), num_filters // 2, name=name + '_for{0}_1'.format(i)), DarknetConv2D_BN_Mish( (3, 3), num_filters // 2 if all_narrow else num_filters, name=name + '_for{0}_2'.format(i))), mode='add')), DarknetConv2D_BN_Mish( (1, 1), num_filters // 2 if all_narrow else num_filters, name=name + '_postconv')) }, mode='concate', name=name + '_route'), DarknetConv2D_BN_Mish((1, 1), num_filters, name=name + '_convblock5'))
def inverted_residual(in_filters, num_filters=64, strides=1, expansion=4, name=''): mid_filters = int(round(in_filters * expansion)) layers = [] if expansion != 1: layers.append( Conv2d_Block((1, 1), num_filters=mid_filters, strides=1, auto_pad=True, padding_mode='zero', normalization='batch', activation='relu6')) layers.append( DepthwiseConv2d_Block((3, 3), depth_multiplier=1, strides=strides, auto_pad=True, padding_mode='zero', normalization='batch', activation='relu6')) layers.append( Conv2d_Block((1, 1), num_filters=num_filters, strides=1, auto_pad=False, padding_mode='zero', normalization='batch', activation=None)) if strides == 1 and in_filters == num_filters: return ShortCut2d(Sequential(*layers), Identity(), activation=None) else: return Sequential(*layers)
def yolo4_body(anchors=None, num_classes=80, image_size=608): anchors1 = generate_anchors(76) if anchors is None else anchors[0] anchors2 = generate_anchors(38) if anchors is None else anchors[1] anchors3 = generate_anchors(19) if anchors is None else anchors[2] num_anchors = len(anchors1) """Create YOLO_V4 model CNN body in Pytorch.""" return Sequential( DarknetConv2D_BN_Mish((3, 3), 32, name='first_layer'), resblock_body(64, 1, all_narrow=False, name='block64'), resblock_body(128, 2, name='block128'), resblock_body(256, 8, name='block256'), ShortCut2d( { 1: Sequential( resblock_body(512, 8, name='block512'), ShortCut2d( { 1: Sequential( resblock_body(1024, 4, name='block1024'), DarknetConv2D_BN_Leaky( (1, 1), 512, name='pre_maxpool1'), DarknetConv2D_BN_Leaky( (3, 3), 1024, name='pre_maxpool2'), DarknetConv2D_BN_Leaky( (1, 1), 512, name='pre_maxpool3'), ShortCut2d( MaxPool2d((13, 13), strides=(1, 1), auto_pad=True), MaxPool2d( (9, 9), strides=(1, 1), auto_pad=True), MaxPool2d( (5, 5), strides=(1, 1), auto_pad=True), Identity(), mode='concate'), DarknetConv2D_BN_Leaky( (1, 1), 512, name='pre_y19_1'), DarknetConv2D_BN_Leaky( (3, 3), 1024, name='pre_y19_2'), DarknetConv2D_BN_Leaky( (1, 1), 512, name='y_19', keep_output=True), DarknetConv2D_BN_Leaky( (1, 1), 256, name='pre_y19_upsample'), Upsampling2d(scale_factor=2, name='y19_upsample'), ), 0: DarknetConv2D_BN_Leaky((1, 1), 256) }, mode='concate'), DarknetConv2D_BN_Leaky((1, 1), 256, name='pre_y38_1'), DarknetConv2D_BN_Leaky((3, 3), 512, name='pre_y38_2'), DarknetConv2D_BN_Leaky((1, 1), 256, name='pre_y38_3'), DarknetConv2D_BN_Leaky((3, 3), 512, name='pre_y38_4'), DarknetConv2D_BN_Leaky( (1, 1), 256, name='y_38', keep_output=True), DarknetConv2D_BN_Leaky( (1, 1), 128, name='pre_y_38_upsample'), Upsampling2d(scale_factor=2, name='y_38_upsample'), ), 0: DarknetConv2D_BN_Leaky((1, 1), 128) }, mode='concate'), DarknetConv2D_BN_Leaky((1, 1), 128, name='pre_y76_concate1'), DarknetConv2D_BN_Leaky((3, 3), 256, name='pre_y76_concate2'), DarknetConv2D_BN_Leaky((1, 1), 128, name='pre_y76_concate3'), DarknetConv2D_BN_Leaky((3, 3), 256, name='pre_y76_concate4'), DarknetConv2D_BN_Leaky((1, 1), 128, name='pre_y76_concate5'), ShortCut2d( #y76_output Sequential(DarknetConv2D_BN_Leaky((3, 3), 256, name='pre_y76_output'), DarknetConv2D((1, 1), num_anchors * (num_classes + 5), use_bias=True, name='y76_output'), YoloLayer(anchors=anchors1, num_classes=num_classes, grid_size=76, img_dim=image_size), name='y76_output'), # y38_output Sequential( ShortCut2d(DarknetConv2D_BN_Leaky((3, 3), 256, strides=(2, 2), auto_pad=False, padding=((1, 0), (1, 0)), name='y76_downsample'), branch_from='y_38', mode='concate'), DarknetConv2D_BN_Leaky((1, 1), 256, name='pre_y38_concate1'), DarknetConv2D_BN_Leaky((3, 3), 512, name='pre_y38_concate2'), DarknetConv2D_BN_Leaky((1, 1), 256, name='pre_y38_concate3'), DarknetConv2D_BN_Leaky((3, 3), 512, name='pre_y38_concate4'), DarknetConv2D_BN_Leaky((1, 1), 256, name='pre_y38_concate5'), ShortCut2d(Sequential(DarknetConv2D_BN_Leaky( (3, 3), 512, name='pre_y38_output'), DarknetConv2D( (1, 1), num_anchors * (num_classes + 5), use_bias=True, name='y38_output'), YoloLayer(anchors=anchors2, num_classes=num_classes, grid_size=38, img_dim=image_size), name='y38_output'), Sequential( ShortCut2d(DarknetConv2D_BN_Leaky( (3, 3), 512, strides=(2, 2), auto_pad=False, padding=((1, 0), (1, 0)), name='y38_downsample'), branch_from='y_19', mode='concate'), DarknetConv2D_BN_Leaky((1, 1), 512, name='pre_y19_concate1'), DarknetConv2D_BN_Leaky((3, 3), 1024, name='pre_y19_concate2'), DarknetConv2D_BN_Leaky((1, 1), 512, name='pre_y19_concate3'), DarknetConv2D_BN_Leaky((3, 3), 1024, name='pre_y19_concate4'), DarknetConv2D_BN_Leaky((1, 1), 512, name='pre_y19_concate5'), Sequential(DarknetConv2D_BN_Leaky( (3, 3), 1024, name='pre_y19_output'), DarknetConv2D( (1, 1), num_anchors * (num_classes + 5), use_bias=True, name='y19_output'), YoloLayer(anchors=anchors3, num_classes=num_classes, grid_size=19, img_dim=image_size), name='y19_output')), mode='concate')), mode='concate'))
def efficient_block(expand_ratio=1, filters_in=32, filters_out=16, kernel_size=3, strides=1, zero_pad=0, se_ratio=0, drop_rate=0.2, is_shortcut=True, name='', **kwargs): expand_ratio = kwargs.get('expand_ratio', expand_ratio) is_shortcut = kwargs.get('id_skip', is_shortcut) filters_in = kwargs.get('filters_in', filters_in) filters_out = kwargs.get('filters_out', filters_out) kernel_size = kwargs.get('kernel_size', kernel_size) is_shortcut = filters_in == filters_out and strides == 1 and kwargs.get( 'id_skip', is_shortcut) filters = filters_in * expand_ratio if expand_ratio == 1 and strides == 1: bottleneck = Sequential( DepthwiseConv2d_Block((kernel_size, kernel_size), depth_multiplier=1, strides=strides, auto_pad=True, padding_mode='zero', normalization='batch', activation='swish'), SqueezeExcite(se_filters=max(1, int(filters_in * se_ratio)), num_filters=filters_in, use_bias=True) if 0 < se_ratio <= 1 else Identity(), Conv2d_Block((1, 1), num_filters=filters_out, strides=1, auto_pad=True, normalization='batch', activation=None), Dropout(dropout_rate=drop_rate) if is_shortcut and drop_rate > 0 else Identity()) if is_shortcut: return ShortCut2d(Identity(), bottleneck) else: return bottleneck else: bottleneck = Sequential( Conv2d_Block((1, 1), num_filters=filters, strides=1, auto_pad=True, normalization='batch', activation='swish'), DepthwiseConv2d_Block((kernel_size, kernel_size), depth_multiplier=1, strides=strides, auto_pad=True, padding_mode='zero', normalization='batch', activation='swish'), SqueezeExcite(se_filters=builtins.max(1, int(filters_in * se_ratio)), num_filters=filters, use_bias=True) if 0 < se_ratio <= 1 else Identity(), Conv2d_Block((1, 1), num_filters=filters_out, strides=1, auto_pad=True, normalization='batch', activation=None), Dropout(dropout_rate=drop_rate) if is_shortcut and drop_rate > 0 else Identity()) if is_shortcut: return ShortCut2d(Identity(), bottleneck) else: return bottleneck
def yolo4_body(num_classes=80, image_size=608): anchors1 = to_tensor(np.array([12, 16, 19, 36, 40, 28]).reshape(-1, 2), requires_grad=False) anchors2 = to_tensor(np.array([36, 75, 76, 55, 72, 146]).reshape(-1, 2), requires_grad=False) anchors3 = to_tensor(np.array([142, 110, 192, 243, 459, 401]).reshape(-1, 2), requires_grad=False) num_anchors = len(anchors1) """Create YOLO_V4 model CNN body in Keras.""" return Sequential( DarknetConv2D_BN_Mish((3, 3), 32, name='first_layer'), resblock_body(64, 1, all_narrow=False, name='block64'), resblock_body(128, 2, name='block128'), resblock_body(256, 8, name='block256'), ShortCut2d( { 1: Sequential( resblock_body(512, 8, name='block512'), ShortCut2d( { 1: Sequential( resblock_body(1024, 4, name='block1024'), DarknetConv2D_BN_Leaky( (1, 1), 512, name='pre_maxpool1'), DarknetConv2D_BN_Leaky( (3, 3), 1024, name='pre_maxpool2'), DarknetConv2D_BN_Leaky( (1, 1), 512, name='pre_maxpool3'), ShortCut2d( MaxPool2d((13, 13), strides=(1, 1), auto_pad=True), MaxPool2d( (9, 9), strides=(1, 1), auto_pad=True), MaxPool2d( (5, 5), strides=(1, 1), auto_pad=True), Identity(), mode='concate'), DarknetConv2D_BN_Leaky( (1, 1), 512, name='pre_y19_1'), DarknetConv2D_BN_Leaky( (3, 3), 1024, name='pre_y19_2'), DarknetConv2D_BN_Leaky( (1, 1), 512, name='y_19', keep_output=True), DarknetConv2D_BN_Leaky( (1, 1), 256, name='pre_y19_upsample'), Upsampling2d(scale_factor=2, name='y19_upsample'), ), 0: DarknetConv2D_BN_Leaky((1, 1), 256) }, mode='concate'), DarknetConv2D_BN_Leaky((1, 1), 256, name='pre_y38_1'), DarknetConv2D_BN_Leaky((3, 3), 512, name='pre_y38_2'), DarknetConv2D_BN_Leaky((1, 1), 256, name='pre_y38_3'), DarknetConv2D_BN_Leaky((3, 3), 512, name='pre_y38_4'), DarknetConv2D_BN_Leaky( (1, 1), 256, name='y_38', keep_output=True), DarknetConv2D_BN_Leaky( (1, 1), 128, name='pre_y_38_upsample'), Upsampling2d(scale_factor=2, name='y_38_upsample'), ), 0: DarknetConv2D_BN_Leaky((1, 1), 128) }, mode='concate'), DarknetConv2D_BN_Leaky((1, 1), 128, name='pre_y76_concate1'), DarknetConv2D_BN_Leaky((3, 3), 256, name='pre_y76_concate2'), DarknetConv2D_BN_Leaky((1, 1), 128, name='pre_y76_concate3'), DarknetConv2D_BN_Leaky((3, 3), 256, name='pre_y76_concate4'), DarknetConv2D_BN_Leaky((1, 1), 128, name='pre_y76_concate5'), ShortCut2d( #y76_output Sequential(DarknetConv2D_BN_Leaky((3, 3), 256, name='pre_y76_output'), DarknetConv2D((1, 1), num_anchors * (num_classes + 5), use_bias=True, name='y76_output'), YoloLayer(anchors=anchors1, num_classes=num_classes, grid_size=76, img_dim=image_size), name='y76_output'), # y38_output Sequential( ShortCut2d(DarknetConv2D_BN_Leaky((3, 3), 256, strides=(2, 2), auto_pad=False, padding=((1, 0), (1, 0)), name='y76_downsample'), branch_from='y_38', mode='concate'), DarknetConv2D_BN_Leaky((1, 1), 256, name='pre_y38_concate1'), DarknetConv2D_BN_Leaky((3, 3), 512, name='pre_y38_concate2'), DarknetConv2D_BN_Leaky((1, 1), 256, name='pre_y38_concate3'), DarknetConv2D_BN_Leaky((3, 3), 512, name='pre_y38_concate4'), DarknetConv2D_BN_Leaky((1, 1), 256, name='pre_y38_concate5'), ShortCut2d(Sequential(DarknetConv2D_BN_Leaky( (3, 3), 512, name='pre_y38_output'), DarknetConv2D( (1, 1), num_anchors * (num_classes + 5), use_bias=True, name='y38_output'), YoloLayer(anchors=anchors2, num_classes=num_classes, grid_size=38, img_dim=image_size), name='y38_output'), Sequential( ShortCut2d(DarknetConv2D_BN_Leaky( (3, 3), 512, strides=(2, 2), auto_pad=False, padding=((1, 0), (1, 0)), name='y38_downsample'), branch_from='y_19', mode='concate'), DarknetConv2D_BN_Leaky((1, 1), 512, name='pre_y19_concate1'), DarknetConv2D_BN_Leaky((3, 3), 1024, name='pre_y19_concate2'), DarknetConv2D_BN_Leaky((1, 1), 512, name='pre_y19_concate3'), DarknetConv2D_BN_Leaky((3, 3), 1024, name='pre_y19_concate4'), DarknetConv2D_BN_Leaky((1, 1), 512, name='pre_y19_concate5'), Sequential(DarknetConv2D_BN_Leaky( (3, 3), 1024, name='pre_y19_output'), DarknetConv2D( (1, 1), num_anchors * (num_classes + 5), use_bias=True, name='y19_output'), YoloLayer(anchors=anchors3, num_classes=num_classes, grid_size=19, img_dim=image_size), name='y19_output')), mode='concate')), mode='concate'))
def __init__(self, *args, axis=1, branch_from=None, activation=None, mode='add', name=None, keep_output=False, **kwargs): """ Args: *args (): axis (): branch_from (): activation (): mode (str): 'add' 'dot' 'concate' name (str): keep_output (bool): **kwargs (): """ super(ShortCut2d, self).__init__(name=name) self.activation = get_activation(activation) self.has_identity = False self.mode = mode self.axis = axis self.branch_from = branch_from self.branch_from_uuid = None self.keep_output = keep_output for i in range(len(args)): arg = args[i] if isinstance(arg, (Layer, torch.Tensor, list, dict)): if isinstance(arg, list): arg = Sequential(*arg) elif isinstance(arg, OrderedDict) and len(args) == 1: for k, v in arg.items(): if isinstance(v, Identity): self.has_identity = True self.add_module('Identity', v) else: self.add_module(k, v) elif isinstance(arg, dict) and len(args) == 1: keys = sorted(list(arg.keys())) for k in keys: v = arg[k] if isinstance(v, Identity): self.has_identity = True self.add_module('Identity', v) else: self.add_module(str(k), v) elif isinstance(arg, (dict, OrderedDict)) and len(args) > 1: raise ValueError( 'more than one dict argument is not support.') elif isinstance(arg, Identity): self.has_identity = True self.add_module('Identity', arg) elif isinstance(arg, nn.Module): if len(arg.name) > 0 and arg.name != arg.default_name: self.add_module(arg.name, arg) else: self.add_module('branch{0}'.format(i + 1), arg) else: raise ValueError('{0} is not support.'.format( arg.__class__.__name)) if len( self._modules ) == 1 and self.has_identity == False and self.branch_from is None and mode != 'concate': self.has_identity = True self.add_module('Identity', Identity()) self.to(self.device)
def build_autoencoder(gan_role='discriminator', use_skip_connections=False): layers = OrderedDict() ae = Sequential() normalization = generator_norm if gan_role == 'generator' else discriminator_norm build_block = generator_build_block if gan_role == 'generator' else discriminator_build_block size_change = 4 initial_size = image_width // pow(2, size_change) filter = 32 embedding_filters = 128 layers['head'] = Conv2d((3, 3), 32, strides=1, auto_pad=True, use_bias=False, activation=activation, name='first_layer') dilation = 1 # WIDTH 128=>64=>32=>16=>8=> # FILTERS 32=>64=>64=>96=>96 for i in range(size_change): if i % 2 == 1: filter = filter + 32 if use_dilation: dilation = 2 if i < 2 else 1 if build_block == BuildBlockMode.base.value: layers['downsample_block{0}'.format(i)] = Conv2d_Block( (3, 3), filter, strides=2, auto_pad=True, use_spectral=use_spectral, use_bias=False, activation=activation, normalization=normalization, name='base_block{0}'.format(i)) elif build_block == BuildBlockMode.resnet.value: layers['downsample_block{0}'.format(i)] = Sequential( resnet_block(filter, strides=2, activation=activation, use_spectral=use_spectral, normalization=normalization, name='resnet_block{0}'.format(i))) elif build_block == BuildBlockMode.bottleneck.value: layers['downsample_block{0}'.format(i)] = bottleneck_block( num_filters=filter, strides=2, reduce=4, activation=activation, use_spectral=use_spectral, normalization=normalization, name='bottleneck_block{0}'.format(i)) if use_self_attention and i == 2: layers['self_attention'] = SelfAttention(16, name='self_attention') layers['embeddings'] = Conv2d((1, 1), embedding_filters, strides=1, auto_pad=True, use_bias=False, activation=activation, keep_output=True, name='embeddings') # 16=>32=>64=>128 # WIDTH 128<=64<=32<=16<=8 # FILTERS 32<=64<=64<=96<=96 for i in range(size_change): if i % 2 == 0: filter = filter - 32 if upsample_mode == UpsampleMode.transpose.value: layers['upsampling{0}'.format( size_change - i - 1)] = TransConv2d_Block( (3, 3), num_filters=filter, strides=2, auto_pad=True, use_spectral=use_spectral, use_bias=False, activation=activation, normalization=normalization, dilation=dilation, name='transconv_block{0}'.format(size_change - i - 1)) elif upsample_mode == UpsampleMode.pixel_shuffle.value: layers['upsampling{0}'.format( size_change - i - 1)] = Sequential( Conv2d_Block((3, 3), num_filters=4 * filter, strides=1, auto_pad=True, use_spectral=use_spectral, use_bias=False, activation=activation, normalization=normalization), Upsampling2d(scale_factor=2, mode=upsample_mode, name='{0}{1}'.format( upsample_mode, size_change - i - 1))) else: layers['upsampling{0}'.format( size_change - i - 1)] = Sequential( Conv2d_Block((3, 3), num_filters=filter, strides=1, auto_pad=True, use_spectral=use_spectral, use_bias=False, activation=activation, normalization=normalization), Upsampling2d(scale_factor=2, mode=upsample_mode, name='{0}{1}'.format( upsample_mode, size_change - i - 1))) if build_block == BuildBlockMode.base.value: layers['upsampling_block{0}'.format( size_change - i - 1)] = Conv2d_Block( (3, 3), filter, strides=1, auto_pad=True, use_spectral=use_spectral, use_bias=False, activation=activation, normalization=normalization, dilation=dilation, name='base_block{0}'.format(size_change - i - 1)) elif build_block == BuildBlockMode.resnet.value: layers['upsampling_block{0}'.format( size_change - i - 1)] = Sequential( resnet_block( filter, strides=1, activation=activation, use_spectral=use_spectral, normalization=normalization, dilation=dilation, name='resnet_block{0}'.format(size_change - i - 1))) elif build_block == BuildBlockMode.bottleneck.value: layers['upsampling_block{0}'.format( size_change - i - 1)] = bottleneck_block( filter, strides=1, activation=activation, use_spectral=use_spectral, normalization=normalization, dilation=dilation, name='resnet_block{0}'.format(size_change - i - 1)) elif build_block == BuildBlockMode.separable_resnet.value: layers['upsampling_block{0}'.format(size_change - i - 1)] = ( separable_resnet_block( filter, strides=2, activation=activation, use_spectral=use_spectral, normalization=discriminator_norm, name='resnet_block{0}'.format(size_change - i - 1))) layers['tail'] = Conv2d((3, 3), generated_output_shape[0], strides=1, auto_pad=True, use_bias=False, activation='tanh', name='last_layer') if use_skip_connections: shortcut3 = ShortCut2d(Identity(), Sequential({ 'downsample_block3': layers['downsample_block3'], 'embeddings': layers['embeddings'], 'upsampling3': layers['upsampling3'], }), mode='add') shortcut2dict = OrderedDict() shortcut2dict['downsample_block2'] = layers['downsample_block2'] if use_self_attention: shortcut2dict['self_attention'] = layers['self_attention'] shortcut2dict['shortcut3'] = shortcut3 shortcut2dict['upsampling_block3'] = layers['upsampling_block3'] shortcut2dict['upsampling2'] = layers['upsampling2'] shortcut2 = ShortCut2d(Identity(), Sequential(shortcut2dict), mode='add') shortcut1 = ShortCut2d(Identity(), Sequential({ 'downsample_block1': layers['downsample_block1'], 'shortcut2': shortcut2, 'upsampling_block2': layers['upsampling_block2'], 'upsampling1': layers['upsampling1'] }), mode='add') shortcut0 = ShortCut2d(Identity(), Sequential({ 'downsample_block0': layers['downsample_block0'], 'shortcut1': shortcut1, 'upsampling_block1': layers['upsampling_block1'], 'upsampling0': layers['upsampling0'] }), mode='add') ae = Sequential({ 'head': layers['head'], 'shortcut0': shortcut0, 'upsampling_block0': layers['upsampling_block0'], 'tail': layers['tail'] }) else: ae = Sequential(layers) return ae