Package: torch-modules-compilation
Torch-Modules-CompilationA compilation of implementations of various ML papers, especially in computer vision. This contains some self-implementations and unofficial & official implementations. More to be added.Install$ pip install torch-modules-compilationTable of ContentsModulesBottleneck Residual BlockDepthwise Seperable ConvolutionSAGAN self-attention moduleGlobal-Local Attention ModuleGlobal Context ModuleLFSA Tokenizer and Refinement BlockParameter-Free Channel Attention (PFCA)Patch MergerResBlockUp/Down sample ResBlockResidual MLP BlockResidual MLP Downsampling BlockTransformer Encoder LayerUNet Encoder and DecoderSqueeze-Excitation ModuleToken LearnerTriplet AttentionLicenseModules/BlocksBottleneck Residual BlockYour basic bottleneck residual block in ResNets.
Image from the paper"Deep Residual Learning for Image Recognition"Parametersin_channels(int): number of input channelsbottleneck_channels(int): number of bottleneck channels; usually less than the number of bottleneck channelsdropout(float): dropout rate; performed after every convolutionUsagefromtorch_modules_compilationimportmodulesx=torch.randn(32,256,16,16)# (batch_size, channels, height, width)block=modules.BottleneckResBlock(in_channels=256,bottleneck_channels=64)block(x).shape# (32, 256, 16, 16)Depthwise Seperable ConvolutionA depthwise seperable convolution; consists of a depthwise convolution and a pointwise convolution. Used in MobileNets and used in the paper"MobileNets: Efficient Convolutional Neural Networks for Mobile Vision Applications". Image also from this paper.Parameters:in_channels(int): Number of input channelsout_channels(int): Number of output channelskernel_size(int): Size of depthwise convolution kernelstride(int): Stride of depthwise convolutionUsagefromtorch_modules_compilationimportmodulesx=torch.randn(32,64,16,16)# (batch_size, channels, height, width)block=modules.DepthwiseSepConv(in_channels=64,out_channels=128,kernel_size=3,stride=1,padding=1)block(x).shape# (32, 128, 16, 16)SAGAN self-attention moduleA feature map self-attention module used in SAGAN;"Self-Attention Generative Adversarial Networks". Image also from this paper. This code implementation was copied and modified fromhttps://github.com/rosinality/sagan-pytorch/blob/master/model.py#L82under Apache 2.0 License. Modification removes spectral initalization.Parametersin_channels(int): Number of input channelsUsagefromtorch_modules_compilationimportmodulesx=torch.randn(32,64,16,16)# (batch_size, channels, height, width)block=modules.FeatureMapSelfAttention(in_channels=64)block(x).shape# (32, 64, 16, 16)Global-Local Attention ModuleAn convolutional attention module introduced in the paper"All the attention you need: Global-local, spatial-channel attention for image retrieval.". Image also from this paper.Parametersin_channels(int): number of channels of the input feature mapnum_reduced_channels(int): number of channels that the local and global spatial attention modules will reduce the input feature map. Refer to figures 3 and 5 in the paper.feaure_map_size(int): height/width of the feature map. The height/width of the input feature maps must be at least 7, due to the 7x7 convolution (3x3 dilated conv) in the module.kernel_size(int): scope of the inter-channel attentionUsagefromtorch_modules_compilationimportmodulesx=torch.randn(32,64,16,16)# (batch_size, channels, height, width)block=modules.GLAM(in_channels=64,num_reduced_channels=48,feature_map_size=16,kernel_size=5)# height and width is equal to feature_map_sizeblock(x).shape# (32, 64, 16, 16)Global Context ModuleA sort of self-attention (non-local) block on feature maps. Implementation of"GCNet: Non-local Networks Meet Squeeze-Excitation Networks and Beyond".Parametersinput_channels(int): Number of input channelsUsagefromtorch_modules_compilationimportmodulesx=torch.randn(32,64,16,16)# (batch_size, channels, height, width)block=modules.GlobalContextModule(input_channels=64)block(x).shape# (32, 64, 16, 16)LFSA Tokenizer and Refinement BlockImplementation of the tokenizer in"Learning Token-Based Representation for Image Retrieval"This are two modules: The tokenizer module that converts feature maps from a CNN (in the paper's case, feature maps from a local-feature-self-attention module) and tokenizes them "into L visual tokens". This is used prior to the refinement block as described in the paper. The refinement block "enhance[s] the obtained visual tokens with self-attention and cross-attention."ParametersLFSA Tokenizerin_channels(int): number of input channelsnum_att_maps(int): number of tokens to tokenize the input into; also the number of channels used by the spatial attentionRefinement Blockd_model(int): dimensionality/channels of inputnhead(int): number of attention heads in the transformerdim_feedforward(int): number of hidden dimensions in the feedforward layersdropout(int): dropout rateUsagefromtorch_modules_compilationimportmodulesx=torch.randn(32,64,16,16)# (batch_size, channels, height, width)tokenizer=modules.LFSATokenizer(in_channels=64,num_att_maps=48)refinement_block=modules.RefinementBlock(d_model=64,nhead=2,dim_feedforward=48*4,dropout=0.1)visual_tokens,cnn_output=tokenizer(x)print(visual_tokens.shape)# (32, 48, 64)print(cnn_output.shape)# (32, 16*16, 64)output=refinement_block(visual_tokens,cnn_output)print(output.shape)# (32, 48, 64)Parameter-Free Channel Attention (PFCA)A channel attention module for convolutional feature maps without any trainable parameters. Used in and image from the paper"PARAMETER-FREE CHANNEL ATTENTION FOR IMAGE CLASSIFICATION AND SUPER-RESOLUTION".Parametersfeature_map_size(int): Length/width of the input feature map_lambda(float): A hyperparameter that is added to the variance (default: 1e-4)Usagefromtorch_modules_compilationimportmodulesx=torch.randn(32,64,16,16)# (batch_size, channels, height, width)block=modules.ParameterFreeChannelAttention(feature_map_size=16)block(x).shape# (32, 64, 16, 16)Patch MergerMerges N tokens into M tokens in transformer models. Typically added in-between transformer layers. Introduced in the paper"LEARNING TO MERGE TOKENS IN VISION TRANSFORMERS". Image from this paper. Copied fromlucidrains' repounder the MIT license.Parametersdim(int): dimensionality/channels of the tokensoutput_tokens(int): number of output merged tokensnorm(bool): normalize the input before mergingscale(bool): scale the attention matrix by the square root of dim (for numerical stability)Usagefromtorch_modules_compilationimportmodulesx=torch.randn(32,64,16)# (batch_size, seq_length, channels)block=modules.PatchMerger(dim=16,output_tokens=48,scale=True)block(x).shape# (32, 48, 16)ResBlockYour basic residual block. Used in ResNets. Image from original paper"Deep Residual Learning for Image Recognition"Parametersin_channels(int): number of input channelskernel_size(int): kernel sizedropout(float): dropout rateUsagefromtorch_modules_compilationimportmodulesx=torch.randn(32,64,16,16)# (batch_size, seq_length, channels)block=modules.ResBlock(in_channels=64,kernel_size=3,dropout=0.2)block(x).shape# (32, 64, 16, 16)Up/Down sample ResBlockComposed of several residual blocks and a down/up sampling at the end; adapted fromStable Diffusion's ResnetBlock.Parametersin_channels(int): number of input channelsout_channels(int): number of output channelsnum_groups(int): number of groups for Group Normalizationnum_layers(int): number of residual blocksdropout(float): dropout ratesample(str): One of "down", "up", or "none". For downsampling 2x, use "down". For upsampling 2x, use "up". Use "none" for no down/up sampling.Usagefromtorch_modules_compilationimportmodulesx=torch.randn(32,64,96,96)# (batch_size, channels, height, width)block=modules.ResBlockUpDownSample(in_channels=64,out_channels=128,num_groups=8,num_layers=2,dropout=0.1,sample='down')block(x).shape# (32, 128, 48, 48)Residual MLP BlockAn improvement of standard MLPs along with residual connections. From"Generalizing MLPs With Dropouts, Batch Normalization, and Skip Connections". This implements the residual MLP block (eq. 5 in the paper).Parametersdim(int): number of input dimensionsic_first(bool): normalize and dropout at the startdropout(float): dropout rateUsagefromtorch_modules_compilationimportmodulesx=torch.randn(32,96)# (batch_size, dim)block=modules.ResidualMLP_block(dim=96,ic_first=True,dropout=0.1)block(x).shape# (32, 96)Residual MLP Downsampling BlockAn improvement of standard MLPs along with residual connections. From"Generalizing MLPs With Dropouts, Batch Normalization, and Skip Connections". This implements the residual MLP block (eq. 6 in the paper).Parametersdim(int): number of input dimensionsdownsample_dim(int): number of output dimensionsdropout(float): dropout rateUsagefromtorch_modules_compilationimportmodulesx=torch.randn(32,96)# (batch_size, dim)block=modules.ResidualMLP_downsample(dim=96,downsample_dim=48,dropout=0.1)block(x).shape# (32, 48)Transformer Encoder LayerStandard transformer encoder layer with queries, keys, and values as inputs.Parametersd_model(int): model dimensionalitynhead(int): number of attention headsdim_feedforward(int): number of hidden dimensions in the feedforward layersdropout(float): dropout ratekdim(int, optional): dimensions of the keysvdim(int, optional): dimensions of the valuesUsagefromtorch_modules_compilationimportmodulesqueries=torch.randn(32,20,64)# (batch_size, seq_length, dim)keys=torch.randn(32,19,48)# (batch_size, seq_length, dim)values=torch.randn(32,19,96)# (batch_size, seq_length, dim)block=modules.TransformerEncoderLayer(d_model=64,nhead=8,dim_feedforward=256,dropout=0.2,kdim=48,vdim=96)block(queries,keys,values).shape# (32, 20, 64)UNet Encoder and DecoderStandard UNet implementation. From the paperU-Net: Convolutional Networks for Biomedical Image Segmentation.ParametersUNet Encoderchannels(list): A list containing the number of channels in the encoder. E.g[3, 64, 128, 256]dropout(float): dropout rateUNet Decoderchannels(list of ints): A list containing the number of channels in the encoder. E.g.[256, 128, 64, 3]dropout(float): dropout rateUsagefromtorch_modules_compilationimportmodulesimages=torch.randn(16,3,224,224)# (batch_size, channels, height, width)unet_encoder=modules.UnetEncoder(channels=[3,64,128,256],dropout=0.1)unet_decoder=modules.UnetDecoder(channels=[256,128,64,3],dropout=0.1)encoder_features=unet_encoder(images)output=unet_decoder(encoder_features)print(output.shape)# (16, 64, 224, 224)Squeeze-Excitation ModuleModule that computes channel-wise interactions in a feature map. FromSqueeze-and-Excitation Networks.Parametersin_channels(int): Number of input channelsreduced_channels(int): Number of channels to reduce to in the "squeeze" part of the modulefeature_map_size(int): height/width of the feature mapUsagefromtorch_modules_compilationimportmodulesfeature_maps=torch.randn(16,128,64,64)# (batch_size, channels, height, width)se_module=modules.SEModule(in_channels=128,reduced_channels=32,feature_map_size=64)se_module(feature_maps)# shape (16, 128, 64, 64); same as inputToken LearnerModule designed for reducing and generating visual tokens given a feature map. FromTokenLearner: What Can 8 Learned Tokens Do for Images and Videos?Parametersin_channels(int): Number of input channelsnum_tokens(int): Number of tokens to reduce toUsagefeature_maps=torch.randn(2,16,10,10)# (batch_size, channels, height, width)token_learner=TokenLearner(in_channels=16,num_tokens=50)# reduce tokens from 10*10 to 50token_learner(feature_maps)# shape (2, 50, 16)Triplet AttentionComputes attention in a feature map across all three dimensions (channel and both spatial dims). FromRotate to Attend: Convolutional Triplet Attention Module.Parametersin_channels(int): Number of input channelsheight(int): height of feature mapwidth(int): width of feature mapkernel_size(int): kernel size of the convolutions. Default: 7Usagefeature_maps=torch.randn(2,16,10,10)# (batch_size, channels, height, width)triplet_attention=TripletAttention(in_channels=16,height=10,width=10)triplet_attention(feature_maps)# shape (2, 16, 10, 10); same as inputLicenseUnless specified, some of these modules are licensed under various licenses and/or copied from other repositories, such as MIT and Apache. Take note of these licenses when using these code in your work. The rest are of my own implementation, which is under the MIT license.See this repo's license file