{ "source_file": "0930", "source_platform": "Medium", "author": "Dhruv Matani", "title_listed": "Efficient Image Segmentation Using PyTorch: Part 4", "keywords": "", "domain": "图像处理", "content_type": "深度技术解析", "is_paid_flag": "否", "url": "https://towardsdatascience.com/efficient-image-segmentation-using-pytorch-part-4-6c86da083432/", "publish_date": "2023-06-27", "corpus_stage": "", "notes": "", "id": "d5fe5f752277", "collected_at": "2026-10-02T10:43:37", "title": "Efficient Image Segmentation Using PyTorch: Part 4 | Towards Data Science", "body_markdown": "![Revisit consent button](../images/body-df29beacd7496220cc94.svg)\nWe value your privacy\nWe use cookies to enhance your browsing experience, serve personalised ads or content, and analyse our traffic. By clicking \"Accept All\", you consent to our use of cookies.\nCustomise Reject All Accept All\nCustomise Consent Preferences ![](../images/body-2159bb67f3de8ad870eb.svg)\nWe use cookies to help you navigate efficiently and perform certain functions. You will find detailed information about all cookies under each consent category below.\nThe cookies that are categorised as \"Necessary\" are stored on your browser as they are essential for enabling the basic functionalities of the site. ... Show more\nNecessaryAlways Active\nNecessary cookies are required to enable the basic features of this site, such as providing secure log-in or adjusting your consent preferences. These cookies do not store any personally identifiable data.\n * Cookie\nBCTempID\n * Duration\n10 minutes\n * Description\nNo description available.\n\n\n * Cookie\n__cf_bm\n * Duration\n1 hour\n * Description\nThis cookie, set by Cloudflare, is used to support Cloudflare Bot Management. \n\n\n * Cookie\nAWSALBCORS\n * Duration\n7 days\n * Description\nAmazon Web Services set this cookie for load balancing.\n\n\n * Cookie\n_cfuvid\n * Duration\nsession\n * Description\nCloudflare sets this cookie to track users across sessions to optimize user experience by maintaining session consistency and providing personalized services\n\n\n * Cookie\nli_gc\n * Duration\n6 months\n * Description\nLinkedin set this cookie for storing visitor's consent regarding using cookies for non-essential purposes.\n\n\n * Cookie\n__hssrc\n * Duration\nsession\n * Description\nThis cookie is set by Hubspot whenever it changes the session cookie. The __hssrc cookie set to 1 indicates that the user has restarted the browser, and if the cookie does not exist, it is assumed to be a new session.\n\n\n * Cookie\n__hssc\n * Duration\n1 hour\n * Description\nHubSpot sets this cookie to keep track of sessions and to determine if HubSpot should increment the session number and timestamps in the __hstc cookie.\n\n\n * Cookie\nwpEmojiSettingsSupports\n * Duration\nsession\n * Description\nWordPress sets this cookie when a user interacts with emojis on a WordPress site. It helps determine if the user's browser can display emojis properly.\n\n\n * Cookie\nBCSessionID\n * Duration\n1 year 1 month 4 days\n * Description\nBlueconic sets this cookie as a unique identifier for the BlueConic profile.\n\n\n * Cookie\n_octo\n * Duration\n1 year\n * Description\nNo description available.\n\n\n * Cookie\nlogged_in\n * Duration\n1 year\n * Description\nNo description available.\n\n\n * Cookie\n__Secure-YEC\n * Duration\npast\n * Description\nYouTube sets this cookie to stores the user's video player preferences using embedded YouTube video\n\n\n * Cookie\n__eoi\n * Duration\n6 months\n * Description\nDescription is currently not available.\n\n\n * Cookie\nAWSALBTGCORS\n * Duration\n7 days\n * Description\nNo description available.\n\n\n * Cookie\nlogin-status-p\n * Duration\npast\n * Description\nDescription is currently not available.\n\n\n * Cookie\nAWSALBTG\n * Duration\n7 days\n * Description\nNo description available.\n\n\n * Cookie\ncsrf_token\n * Duration\nsession\n * Description\nNo description available.\n\n\n * Cookie\ntoken_v2\n * Duration\n1 day\n * Description\nDescription is currently not available.\n\n\n * Cookie\nD\n * Duration\n1 year\n * Description\nDescription is currently not available.\n\n\n * Cookie\nPHPSESSID\n * Duration\nsession\n * Description\nThis cookie is native to PHP applications. The cookie stores and identifies a user's unique session ID to manage user sessions on the website. The cookie is a session cookie and will be deleted when all the browser windows are closed.\n\n\n * Cookie\nVISITOR_PRIVACY_METADATA\n * Duration\n6 months\n * Description\nYouTube sets this cookie to store the user's cookie consent state for the current domain. \n\n\n * Cookie\ncookietest\n * Duration\nsession\n * Description\nThe cookietest cookie is typically used to determine whether the user's browser accepts cookies, essential for website functionality and user experience.\n\n\n * Cookie\n__Host-airtable-session\n * Duration\n1 year\n * Description\nThis cookie is used to enable us to integrate the services of Airtable.\n\n\n * Cookie\n__Host-airtable-session.sig\n * Duration\n1 year\n * Description\nThis cookie is used to enable us to integrate the services of Airtable.\n\n\n * Cookie\nm\n * Duration\n1 year 1 month 4 days\n * Description\nStripe sets this cookie for fraud prevention purposes. It identifies the device used to access the website, allowing the website to be formatted accordingly.\n\n\n * Cookie\nBIGipServer*\n * Duration\nsession\n * Description\nMarketo sets this cookie to collect information about the user's online activity and build a profile about their interests to provide advertisements relevant to the user. \n\n\n * Cookie\n__cfruid\n * Duration\nsession\n * Description\nCloudflare sets this cookie to identify trusted web traffic.\n\n\n * Cookie\n_GRECAPTCHA\n * Duration\n6 months\n * Description\nGoogle Recaptcha service sets this cookie to identify bots to protect the website against malicious spam attacks.\n\n\n * Cookie\n__Secure-YNID\n * Duration\n6 months\n * Description\nGoogle cookie used to protect user security and prevent fraud, especially during the login process.\n\n\n * Cookie\ncookieyes-consent\n * Duration\n1 year\n * Description\nCookieYes sets this cookie to remember users' consent preferences so that their preferences are respected on subsequent visits to this site. It does not collect or store any personal information about the site visitors.\n\n\nFunctional\nFunctional cookies help perform certain functionalities like sharing the content of the website on social media platforms, collecting feedback, and other third-party features.\n * Cookie\nlidc\n * Duration\n1 day\n * Description\nLinkedIn sets the lidc cookie to facilitate data center selection.\n\n\n * Cookie\nbrw\n * Duration\n1 year\n * Description\nNo description available.\n\n\n * Cookie\nbrwConsent\n * Duration\n5 minutes\n * Description\nDescription is currently not available.\n\n\n * Cookie\nWMF-Uniq\n * Duration\n1 year\n * Description\nDescription is currently not available.\n\n\n * Cookie\nloom_anon_comment\n * Duration\n1 year\n * Description\nNo description available.\n\n\n * Cookie\nloom_referral_video\n * Duration\nsession\n * Description\nDescription is currently not available.\n\n\n * Cookie\nVISITOR_INFO1_LIVE\n * Duration\n6 months\n * Description\nA cookie set by YouTube to measure bandwidth that determines whether the user gets the new or old player interface.\n\n\n * Cookie\nyt-remote-connected-devices\n * Duration\nNever Expires\n * Description\nYouTube sets this cookie to store the user's video preferences using embedded YouTube videos.\n\n\n * Cookie\nytidb::LAST_RESULT_ENTRY_KEY\n * Duration\nNever Expires\n * Description\nThe cookie ytidb::LAST_RESULT_ENTRY_KEY is used by YouTube to store the last search result entry that was clicked by the user. This information is used to improve the user experience by providing more relevant search results in the future.\n\n\n * Cookie\nyt-remote-device-id\n * Duration\nNever Expires\n * Description\nYouTube sets this cookie to store the user's video preferences using embedded YouTube videos.\n\n\n * Cookie\nyt-remote-session-name\n * Duration\nsession\n * Description\nThe yt-remote-session-name cookie is used by YouTube to store the user's video player preferences using embedded YouTube video.\n\n\n * Cookie\nyt-remote-fast-check-period\n * Duration\nsession\n * Description\nThe yt-remote-fast-check-period cookie is used by YouTube to store the user's video player preferences for embedded YouTube videos.\n\n\n * Cookie\nyt-remote-session-app\n * Duration\nsession\n * Description\nThe yt-remote-session-app cookie is used by YouTube to store user preferences and information about the interface of the embedded YouTube video player.\n\n\n * Cookie\nyt-remote-cast-available\n * Duration\nsession\n * Description\nThe yt-remote-cast-available cookie is used to store the user's preferences regarding whether casting is available on their YouTube video player.\n\n\n * Cookie\nyt-remote-cast-installed\n * Duration\nsession\n * Description\nThe yt-remote-cast-installed cookie is used to store the user's video player preferences using embedded YouTube video.\n\n\n * Cookie\ncp_session\n * Duration\n3 months\n * Description\nCodepen sets this cookie for Help systems found in the website.\n\n\n * Cookie\nloid\n * Duration\n1 year 1 month 4 days\n * Description\nThis cookie is set by the Reddit. The cookie enables the sharing of content from the website onto the social media platform.\n\n\nAnalytics\nAnalytical cookies are used to understand how visitors interact with the website. These cookies help provide information on metrics such as the number of visitors, bounce rate, traffic source, etc.\n * Cookie\n__hstc\n * Duration\n6 months\n * Description\nHubspot set this main cookie for tracking visitors. It contains the domain, initial timestamp (first visit), last timestamp (last visit), current timestamp (this visit), and session number (increments for each subsequent session).\n\n\n * Cookie\nhubspotutk\n * Duration\n6 months\n * Description\nHubSpot sets this cookie to keep track of the visitors to the website. This cookie is passed to HubSpot on form submission and used when deduplicating contacts.\n\n\n * Cookie\n_ga\n * Duration\n1 year 1 month 4 days\n * Description\nGoogle Analytics sets this cookie to calculate visitor, session and campaign data and track site usage for the site's analytics report. The cookie stores information anonymously and assigns a randomly generated number to recognise unique visitors.\n\n\n * Cookie\n_ga_*\n * Duration\n1 year 1 month 4 days\n * Description\nGoogle Analytics sets this cookie to store and count page views.\n\n\n * Cookie\n__Host-psifi.analyticsTrace\n * Duration\n6 hours\n * Description\nDescription is currently not available.\n\n\n * Cookie\n__Host-psifi.analyticsTraceV2\n * Duration\n6 hours\n * Description\nDescription is currently not available.\n\n\n * Cookie\n_gh_sess\n * Duration\nsession\n * Description\nGitHub sets this cookie for temporary application and framework state between pages like what step the user is on in a multiple step form.\n\n\n * Cookie\nYSC\n * Duration\nsession\n * Description\nYSC cookie is set by Youtube and is used to track the views of embedded videos on Youtube pages.\n\n\n * Cookie\najs_anonymous_id\n * Duration\n1 year\n * Description\nThis cookie is set by Segment to count the number of people who visit a certain site by tracking if they have visited before.\n\n\n * Cookie\nvuid\n * Duration\n1 year 1 month 4 days\n * Description\nVimeo installs this cookie to collect tracking information by setting a unique ID to embed videos on the website. \n\n\nPerformance\nPerformance cookies are used to understand and analyse the key performance indexes of the website which helps in delivering a better user experience for the visitors.\n * Cookie\nAWSALB\n * Duration\n7 days\n * Description\nAWSALB is an application load balancer cookie set by Amazon Web Services to map the session to the target.\n\n\n * Cookie\nacq\n * Duration\npast\n * Description\nDescription is currently not available.\n\n\n * Cookie\nacq.sig\n * Duration\npast\n * Description\nDescription is currently not available.\n\n\n * Cookie\nptc\n * Duration\n2 years\n * Description\nNo description available.\n\n\nAdvertisement\nAdvertisement cookies are used to provide visitors with customised advertisements based on the pages you visited previously and to analyse the effectiveness of the ad campaigns.\n * Cookie\nmuc_ads\n * Duration\n1 year 1 month 4 days\n * Description\nTwitter sets this cookie to collect user behaviour and interaction data to optimize the website.\n\n\n * Cookie\nguest_id_marketing\n * Duration\n1 year 1 month 4 days\n * Description\nTwitter sets this cookie to identify and track the website visitor.\n\n\n * Cookie\nguest_id_ads\n * Duration\n1 year 1 month 4 days\n * Description\nTwitter sets this cookie to identify and track the website visitor.\n\n\n * Cookie\npersonalization_id\n * Duration\n1 year 1 month 4 days\n * Description\nTwitter sets this cookie to integrate and share features for social media and also store information about how the user uses the website, for tracking and targeting.\n\n\n * Cookie\nguest_id\n * Duration\n1 year 1 month 4 days\n * Description\nTwitter sets this cookie to identify and track the website visitor. It registers if a user is signed in to the Twitter platform and collects information about ad preferences.\n\n\n * Cookie\nbcookie\n * Duration\n1 year\n * Description\nLinkedIn sets this cookie from LinkedIn share buttons and ad tags to recognize browser IDs.\n\n\n * Cookie\n__Secure-ROLLOUT_TOKEN\n * Duration\n6 months\n * Description\nYouTube sets this cookie to manage feature rollout and experimentation. It helps Google control which new features or interface changes are shown to users as part of testing and staged rollouts, ensuring consistent experience for a given user during an experiment.\n\n\n * Cookie\nyt.innertube::nextId\n * Duration\nNever Expires\n * Description\nYouTube sets this cookie to register a unique ID to store data on what videos from YouTube the user has seen.\n\n\n * Cookie\nyt.innertube::requests\n * Duration\nNever Expires\n * Description\nYouTube sets this cookie to register a unique ID to store data on what videos from YouTube the user has seen.\n\n\n * Cookie\nsession_tracker\n * Duration\nsession\n * Description\nThis cookie is set by the Reddit. This cookie is used to identify trusted web traffic. It also helps in adverstising on the website.\n\n\n * Cookie\nedgebucket\n * Duration\nsession\n * Description\nReddit sets this cookie to save the information about a log-on Reddit user, for the purpose of advertisement recommendations and updating the content.\n\n\n * Cookie\ndid\n * Duration\n1 year\n * Description\nArbor sets this cookie to show targeted ads to site visitors.This cookie expires after 2 months or 1 year.\n\n\nUncategorised\nOther uncategorised cookies are those that are being analysed and have not been classified into a category as yet.\nNo cookies to display.\nReject All Save My Preferences Accept All\nToggle Search[](https://towardsdatascience.com/)Toggle Mobile Navigation\n[](https://towardsdatascience.com/)\nPublish AI, ML & data-science insights to a global community of data professionals.\nSign in[Submit an Article](https://write.insightmediagroup.io)\n[Latest](https://towardsdatascience.com/latest/)[Editor's Picks](https://towardsdatascience.com/tag/editors-pick/)[Newsletter](https://towardsdatascience.com/tag/the-variable/)[ShipAI](https://towardsdatascience.com/shipai/)\n* * *\n[Write For TDS](https://towardsdatascience.com/submissions/)\n[LinkedIn](https://www.linkedin.com/company/towards-data-science/)[X](https://x.com/TDataScience)Toggle Search\n# Efficient Image Segmentation Using PyTorch: Part 4\nA Vision Transformer-based model\n[Dhruv Matani](https://towardsdatascience.com/author/dhruvbird/)\nJune 27, 202315 min read\nIn this 4-part series, we'll implement image segmentation step by step from scratch using deep learning techniques in PyTorch. This part will focus on implementing a Vision Transformer based model for image segmentation.\nCo-authored with [Naresh Singh](https://medium.com/@brocolishbroxoli)\n![Figure 1: Result of running image segmentation using a vision transformer model architecture. From top to bottom, input images, ground truth segmentation masks, and predicted segmentation masks. Source: Author\\(s\\)](../images/000.png)\nFigure 1: Result of running image segmentation using a vision transformer model architecture. From top to bottom, input images, ground truth segmentation masks, and predicted segmentation masks. Source: Author(s)\n* * *\n## Article outline\nIn this article, we will visit the [Transformer architecture](https://ai.googleblog.com/2017/08/transformer-novel-neural-network.html) which has taken the world of deep learning by storm. The transformer is a multimodal architecture which can model different modalities such as language, vision, and audio.\nIn this article, we will\n 1. Learn about the transformer architecture and the key concepts involved\n 2. Understand the vision transformer architecture\n 3. Introduce a vision transformer model that is written from scratch so that you can appreciate all the building blocks and moving parts\n 4. Follow an input tensor fed into this model and inspect how it changes shape\n 5. Use this model to perform image segmentation on the Oxford IIIT Pet dataset\n 6. Observe the results of this segmentation task\n 7. Briefly introduce the SegFormer, a state of the art vision transformer for semantic segmentation\n\n\nThroughout this article, we will reference code and results from this [notebook](https://github.com/dhruvbird/ml-notebooks/blob/main/pets_segmentation/vision-transformer-for-oxford-iiit-pet-segmentatio.ipynb) for model training. If you wish to reproduce the results, you'll need a GPU to ensure that the first notebook completes running in a reasonable amount of time.\n* * *\n## Articles in this series\nThis series is for readers at all experience levels with deep learning. If you want to learn about the practice of deep learning and vision AI along with some solid theory and hands-on experience, you've come to the right place! This is expected to be a 4-part series with the following articles:\n 1. [Concepts and Ideas](https://medium.com/p/89e8297a0923/)\n 2. [A CNN-based model](https://medium.com/p/bed68cadd7c7/)\n 3. [Depthwise separable convolutions](https://medium.com/p/3534cf04fb89/)\n 4. **A Vision Transformer-based model (this article)**\n\n\nLet's start our journey into vision transformers with an introduction and intuitive understanding of the transformer architecture.\n* * *\n## The Transformer Architecture\nWe can think of the transformer architecture as a composition of interleaving layers of _communication_ and _computation_. This idea is depicted visually in figure 2. The transformer has N processing units (N is 3 in Figure 2), each of which is responsible for processing a 1/N fraction of the input. For those processing units to produce meaningful results, each of them need to have a _global view_ of the input. Hence, the system repeatedly communicates information about the data in every processing unit to every other processing unit; shown using the red, green, and blue arrows going from every processing unit to every other processing unit. This is followed by some computation based on this information. After sufficient repetitions of this process, the model is able to produce desired results.\n![Figure 2: Interleaved communication and computation in transformers. The image shows just 2 layers of communication and computation. In practice, there are many more such layers. Source: Author\\(s\\).](../images/001.png)\nFigure 2: Interleaved communication and computation in transformers. The image shows just 2 layers of communication and computation. In practice, there are many more such layers. Source: Author(s).\nIt's worth noting most online resources typically discuss both the encoder and the decoder of the transformer as presented in the paper titled \"[Attention is all you need](https://arxiv.org/abs/1706.03762).\" However, in this article, we will describe just the encoder part of the transformer.\nLet's take a closer look at what constitutes communication and computation in transformers.\n### Communication in transformers: Attention\nIn transformers, communication is implemented by a layer known as the attention layer. In PyTorch, this is called [MultiHeadAttention](https://pytorch.org/docs/stable/generated/torch.nn.MultiheadAttention.html). We'll get to the reason for that name in a bit.\nThe documentation says:\n_\"Allows the model to jointly attend to information from different representation subspaces as described in the paper:_[_Attention is all you need_](https://arxiv.org/abs/1706.03762) _.\"_\nThe attention mechanism consumes an input tensor _x_ of shape (Batch, Length, Features), and it produces a similarly shaped tensor _y_ such that the features for each input are updated based on which other inputs in the same instance the tensor is paying attention to. Hence, the features of each tensor of length \"Features\" in the instance of size \"Length\" are updated based on every other tensor. This is where the quadratic cost of the attention mechanism comes in.\n![Figure 3: Attention of the word "it " shown relative to the other words in the sentence. We can see that "it " is paying attention to the words "animal ", "too ", and "tire\\(d\\) " in the same sentence. Source: Generated using this colab.](../images/002.png)\nFigure 3: Attention of the word \"it \" shown relative to the other words in the sentence. We can see that \"it \" is paying attention to the words \"animal \", \"too \", and \"tire(d) \" in the same sentence. Source: Generated using [this colab](https://colab.research.google.com/github/tensorflow/tensor2tensor/blob/master/tensor2tensor/notebooks/hello_t2t.ipynb).\nIn the context of a vision transformer, the input to the transformer is an image. Let's assume this to be a 128 x 128 (width, height) image. We chunk it into multiple smaller patches of size (16 x 16). For a 128 x 128 image, we get 64 patches (Length), 8 patches in each row, and 8 rows of patches.\nEach one of these 64 patches of size 16 x 16 pixels is considered to be a separate input to the transformer model. Without getting too deep into the details, it should be sufficient to think of this process as being driven by 64 different processing units, each of which is processing a single 16x16 image patch.\nIn each round, the attention mechanism in each processing unit is responsible for looking at the image patch it is responsible for and querying each one of the other remaining 63 processing units to ask them for any information that may be relevant and useful to help it effectively process its own image patch.\nThe communication step via attention is followed by computation, which we will look at next.\n### Computation in transformers: Multi Layer Perceptron\nComputation in transformers is nothing but a MultiLayerPerceptron (MLP) unit. This unit is composed of 2 Linear layers, with a GeLU non-linearity in between. One can consider using other non-linearities as well. This unit first projects the input to 4x the size and reprojects it back to 1x, which is the same as the input size.\nIn the code we'll see in our notebook, this class is called MultiLayerPerceptron. The code is shown below.\npython\n\n```\nclass MultiLayerPerceptron(nn.Sequential):\n def __init__(self, embed_size, dropout):\n super().__init__(\n nn.Linear(embed_size, embed_size * 4),\n nn.GELU(),\n nn.Linear(embed_size * 4, embed_size),\n nn.Dropout(p=dropout),\n )\n # end def\n# end class\n```\n\nNow that we understand the high level working of the transformer architecture, let's focus our attention on the vision transformer since we're going to be performing image segmentation.\n* * *\n## The Vision Transformer\nThe vision transformer was first introduced by the paper titled \"[An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://arxiv.org/abs/2010.11929)\". The paper discusses how the authors apply the vanilla transformer architecture to the problem of image classification. This is done by splitting the image into patches of size 16x16, and treating each patch as an input token to the model. The transformer encoder model is fed these input tokens, and is asked to predict a class for the input image.\n![Figure 4: Source: Transformers for image recognition at scale.](../images/003.png)\nFigure 4: Source: [Transformers for image recognition at scale](https://arxiv.org/pdf/2010.11929.pdf).\nIn our case, we are interested in image segmentation. We can consider it to be a pixel-level classification task because we intend to predict a target class per pixel..\nWe make a small but important change to the vanilla vision transformer and replace the MLP head for classification by an MLP head for pixel level classification. We have a single linear layer in the output that is shared by every patch whose segmentation mask is predicted by the vision transformer. This shared linear layer predicts a segmentation mask for every patch that was sent as input to the model.\nIn the case of the vision transformer, a patch of size 16x16 is considered to be equivalent to a single input token at a specific time step.\n![Figure 5: The end to end working of the vision transformer for image segmentation. Image generated using this notebook. Source: Author\\(s\\).](../images/004.png)\nFigure 5: The end to end working of the vision transformer for image segmentation. Image generated using this [notebook](https://github.com/dhruvbird/ml-notebooks/blob/main/pets_segmentation/Vision%20Transformer%20Article%20Images.ipynb). Source: Author(s).\n### Building an intuition for tensor dimensions in vision transformers\nWhen working with deep CNNs, the tensor dimensions we used for the most part was (N, C H, W), where the letters stand for the following:\n * N: Batch size\n * C: Number of channels\n * H: Height\n * W: Width\n\n\nYou can see that this format is geared toward 2d image processing, since it smells of features that are very specific to images.\nWith transformers on the other hand, things become a lot more generic and domain agnostic. What we'll see below applies to vision, text, NLP, audio or other problems where input data can be represented as a sequence. It is worth noting that there's little vision specific bias in the representation of tensors as they flow through our vision transformer.\nWhen working with transformers and attention in general, we expect the tensors to have the following shape: (B, T, C), where the letters stand for the following:\n * B: Batch size (same as that for CNNs)\n * T: Time dimension or sequence length. This dimension is also sometimes called L. In the case of vision transformers, each image patch corresponds to this dimension. If we have 16 image patches, then the value of the T dimension will be 16\n * C: The channel or embedding size dimension. This dimension is also sometimes called E. When processing images, each patch of size 3x16x16 (Channel, Width, Height) is mapped via a patch embedding layer to an embedding of size C. We'll see how this is done later.\n\n\nLet's dive into how the input image tensor gets mutated and processed along its way to predicting the segmentation mask.\n### The journey of a tensor in a vision transformer\nIn deep CNNs, the journey of a tensor looks something like this (in a UNet, SegNet, or other CNN based architecture).\nThe input tensor is typically of shape (1, 3, 128, 128). This tensor goes through a series of convolution and max-pooling operations where its spatial dimensions are reduced and channel dimensions are increased, typically by a factor of 2 each. This is called the feature encoder. After this, we do the reverse operation where we increase the spatial dimensions and reduce the channel dimensions. This is called the feature decoder. After the decoding process, we get a tensor of shape (1, 64, 128, 128). This is then projected into the number of output channels C that we desire as (1, C, 128, 128) using a 1x1 pointwise convolution without bias.\n![Figure 6: Typical progression of tensor shapes through a deep CNN used for image segmentation. Source: Author\\(s\\).](../images/005.png)\nFigure 6: Typical progression of tensor shapes through a deep CNN used for image segmentation. Source: Author(s).\nWith vision transformers, the flow is much more complex. Let's take a look at an image below and then try to understand how the tensor transforms shapes at every step along the way.\n![Figure 7: Typical progression of tensor shapes through a vision transformer for image segmentation. Source: Author\\(s\\).](../images/006.png)\nFigure 7: Typical progression of tensor shapes through a vision transformer for image segmentation. Source: Author(s).\nLet's look at each step in more detail and see how it updates the shape of the tensor flowing through the vision transformer. To understand this better, let's take concrete values for our tensor dimensions.\n 1. **Batch Normalization:** The input and output tensors have shape (1, 3, 128, 128). The shape is unchanged, but the values are normalized to zero mean and unit variance.\n 2. **Image to patches:** The input tensor of shape (1, 3, 128, 128) is converted into a stacked patch of 16x16 images. The output tensor has shape (1, 64, 768).\n 3. **Patch embedding:** The patch embedding layer maps the 768 input channels to 512 embedding channels (for this example). The output tensor is of shape (1, 64, 512). The patch embedding layer is basically just an nn.Linear layer in PyTorch.\n 4. **Position embedding:** The position embedding layer doesn't have an input tensor, but effectively contributes a learnable parameter (trainable tensor in PyTorch) o f the same shape as the patch embedding. This is of shape (1, 64, 512).\n 5. **Add:** The patch and position embeddings are added together piecewise to produce the input to our vision transformer encoder. This tensor is of shape (1, 64, 512). You'll notice that the main workhorse of the vision transformer, i.e. the encoder basically leaves this tensor shape unchanged.\n 6. **Transformer encoder:** The input tensor of shape (1, 64, 512) flows through multiple transformer encoder blocks, each of which have multiple attention heads (communication) followed by an MLP layer (computation). The tensor shape remains unchanged as (1, 64, 512).\n 7. **Linear output projection:** If we assume that we want to segment each image into 10 classes, then we will need each patch of size 16x16 to have 10 channels. The nn.Linear layer for output projection will now convert the 512 embedding channels to 16x16x10 = 2560 output channels, and this tensor will look like (1, 64, 2560). In the diagram above C' = 10. Ideally, this would be a multi-layer perceptron, since _\"[MLPs are universal function approximators](https://en.wikipedia.org/wiki/Multilayer_perceptron)\"_, but we use a single linear layer since this is an educational exercise\n 8. **Patch to image:** This layer converts the 64 patches encoded as a (1, 64, 2560) tensor back into something that looks like a segmentation mask. This can be 10 single channel images, or in this case a single 10 channel image, with each channel being the segmentation mask for one of the 10 classes. The output tensor is of shape (1, 10, 128, 128).\n\n\nThat's it - we've successfully segmented an input image using a vision transformer! Next, let's take a look at an experiment along with some results.\n### Vision transformers in action\nThis [notebook](https://github.com/dhruvbird/ml-notebooks/blob/main/pets_segmentation/vision-transformer-for-oxford-iiit-pet-segmentatio.ipynb) contains all the code for this section.\nAs far as the code and class structure is concerned, it closely mimics the block diagram above. Most of the concepts mentioned above have a 1:1 correspondence to class names in this [notebook](https://github.com/dhruvbird/ml-notebooks/blob/main/pets_segmentation/vision-transformer-for-oxford-iiit-pet-segmentatio.ipynb).\nThere are some concepts related to the attention layers that are critical hyperparameters for our model. We didn't mention anything about the details of the multi-head attention earlier since we mentioned that it's out of scope for the purposes of this article. We highly recommend reading the reference material mentioned above before proceeding if you don't have a basic understanding of the attention mechanism in transformers.\nWe used the following model parameters for the vision transformer for segmentation.\n 1. 768 embedding dimensions for the PatchEmbedding layer\n 2. 12 Transformer encoder blocks\n 3. 8 attention heads in each transformer encoder block\n 4. 20% dropout in multi-head attention and MLP\n\n\nThis configuration can be seen in the VisionTransformerArgs Python dataclass.\npython\n\n```\n@dataclass\nclass VisionTransformerArgs:\n \"\"\"Arguments to the VisionTransformerForSegmentation.\"\"\"\n image_size: int = 128\n patch_size: int = 16\n in_channels: int = 3\n out_channels: int = 3\n embed_size: int = 768\n num_blocks: int = 12\n num_heads: int = 8\n dropout: float = 0.2\n# end class\n```\n\nA similar configuration as [before](https://medium.com/p/bed68cadd7c7/) was used during model training and validation. The configuration is specified below.\n 1. The _random horizontal flip_ and _colour jitter_ data augmentations are applied to the training set to prevent overfitting\n 2. The images are resized to 128x128 pixels in a non-aspect preserving resize operation\n 3. No input normalization is applied to the images - instead a [batch normalization layer is used as the first layer of the model](https://towardsdatascience.com/replace-manual-normalization-with-batch-normalization-in-vision-ai-models-e7782e82193c)\n 4. The model is trained for 50 epochs using the Adam optimizer with a LR of 0.0004 and a StepLR scheduler that decays the learning rate by 0.8x every 12 epochs\n 5. The cross-entropy loss function is used to classify a pixel as belonging to a pet, the background, or a pet border\n\n\nThe model has 86.28M parameters and achieved a validation accuracy of 85.89% after 50 training epochs. This is less than the 88.28% accuracy achieved by deep CNN model after 20 training epochs. This could be due to a few factors that need to be validated experimentally.\n 1. The last output projection layer is a single nn.Linear and not a multi-layer perceptron\n 2. The 16x16 patch size is too large to capture more fine grained detail\n 3. Not enough training epochs\n 4. Not enough training data - it's known that transformer models need a lot more data to train effectively compared to deep CNN models\n 5. The learning rate is too low\n\n\nWe plotted a gif showing how the model is learning to predict the segmentation masks for 21 images in the validation set.\n![Figure 8: A gif showing the progression of segmentation masks predicted by the vision transformer for image segmentation model. Source: Author\\(s\\).](../images/body-013c1b7c56b4d349ea0f.img)\nFigure 8: A gif showing the progression of segmentation masks predicted by the vision transformer for image segmentation model. Source: Author(s).\nWe notice something interesting in the early training epochs. The predicted segmentation masks have some strange blocking artifacts. The only reason we could think of for this is because we're breaking down the image into patches of size 16x16 and after very few training epochs, the model hasn't learned anything useful beyond some very coarse grained information regarding whether this 16x16 patch is generally covered by a pet or by background pixels.\n![Figure 9: The blocking artifacts seen in the predicted segmentation masks when using the vision transformer for image segmentation. Source: Author\\(s\\).](../images/007.png)\nFigure 9: The blocking artifacts seen in the predicted segmentation masks when using the vision transformer for image segmentation. Source: Author(s).\nNow that we have seen a basic vision transformer in action, let's turn our attention to a state of the art vision transformer for segmentation tasks.\n### SegFormer: Semantic segmentation with transformers\nThe SegFormer architecture was proposed in [this paper](https://arxiv.org/pdf/2105.15203.pdf) in 2021. The transformer we saw above is a simpler version of the SegFormer architecture.\n![Figure 10: The SegFormer architecture. Source: SegFormer paper \\(2021\\).](../images/008.png)\nFigure 10: The SegFormer architecture. Source: [SegFormer paper (2021)](https://arxiv.org/pdf/2105.15203.pdf).\nMost notably, the SegFormer:\n 1. Generates 4 sets of images with patches of size 4x4, 8x8, 16x16, and 32x32 instead of a single patched image with patches of size 16x16\n 2. Uses 4 transformer encoder blocks instead of just 1. This feels like a model ensemble\n 3. Uses convolutions in the pre and post phases of self-attention\n 4. Doesn't use positional embeddings\n 5. Each transformer block processes images at spatial resolution H/4 x W/4, H/8 x W/8, H/16 x W/16, and H/32, W/32\n 6. Similarly, the channels increase when the spatial dimensions reduce. This feels similar to deep CNNs\n 7. Predictions at multiple spatial dimensions are upsampled and then merged together in the decoder\n 8. An MLP combines all these predictions to provide a final prediction\n 9. The final prediction is at spatial dimension H/4, W/4 and not at H, W.\n\n\n* * *\n## Conclusion\nIn part-4 of this series, we were introduced to the transformer architecture and vision transformers in particular. We developed an intuitive understanding of how vision transformers work, and the basic building block involved in the communication and computation phases of vision transformers. We saw the unique patch based approach adopted by vision transformers for predicting segmentation masks and then combining the predictions together.\nWe reviewed an experiment that shows vision transformers in action, and were able to compare results with deep CNN approaches. While our vision transformer is not state of the art, it was able to achieve pretty decent results. We provided a glimpse into state-of-the-art approaches such as SegFormer.\nIt should be clear by now that transformers have a lot more moving parts and are more complex compared to deep CNN-based approaches. From a raw FLOPs point of view, transformers hold the promise of being more efficient. In transformers, the only real layer that is computationally heavy is nn.Linear. This is implemented using optimized matrix multiplication on most architectures. Due to this architectural simplicity, transformers hold the promise of being easier to optimize and speed up compared to deep CNN-based approaches.\nCongratulations on making it this far! We're glad you enjoyed reading this series on efficient image segmentation in PyTorch. If you have questions or comments, please feel free to leave them in the comments section.\n* * *\n## Further reading\nThe details of the attention mechanism are out of scope for this article. Besides, there are numerous high quality resources that you can refer to to understand the attention mechanism in great detail. Here are some that we highly recommend.\n 1. [The Illustrated Transformer](http://jalammar.github.io/illustrated-transformer/)\n 2. [NanoGPT from scratch using PyTorch](https://www.youtube.com/watch?v=kCc8FmEb1nY)\n\n\nWe'll provide links to articles that provide more details on vision transformers below.\n 1. [Implementing Vision Transformer (ViT) in PyTorch](https://towardsdatascience.com/implementing-visualttransformer-in-pytorch-184f9f16f632): This article details the implementation of a vision transformer for image classification in PyTorch. Notably, their implementation uses [einops](https://einops.rocks/), which we avoid, since this is an education-focused exercise (we recommend learning and using einops for code readability though). We instead use native PyTorch operators for permuting and rearranging tensor dimensions. Additionally, there are a few places where the author uses Conv2d instead of Linear layers. We wanted to build an implementation of vision transformers without the use of convolutional layers entirely.\n 2. [Vision Transformer: AI Summer](https://theaisummer.com/vision-transformer/)\n 3. [Implementing SegFormer in PyTorch](https://towardsdatascience.com/implementing-segformer-in-pytorch-8f4705e2ed0e)\n\n\n## Register Now to Continue Reading\n### Unlock Towards Data Science with a Free Account\nAll articles on Towards Data Science are available to registered users. Create a free account today to access this article and thousands more.\nRegister for freeSign in to your account\n* * *\n![](../images/body-a3403563a948ce1eab7b.jpeg)\nWritten By\n## [Dhruv Matani](https://towardsdatascience.com/author/dhruvbird/)\n[See all from Dhruv Matani](https://towardsdatascience.com/author/dhruvbird/)\n[Deep Dives](https://towardsdatascience.com/tag/deep-dives/)[Pytorch](https://towardsdatascience.com/tag/pytorch/)[Pytorch Image Segment](https://towardsdatascience.com/tag/pytorch-image-segment/)[Transformers](https://towardsdatascience.com/tag/transformers/)[Vision Transformer](https://towardsdatascience.com/tag/vision-transformer/)\nShare This Article\n * [](https://www.facebook.com/sharer/sharer.php?u=https%3A%2F%2Ftowardsdatascience.com%2Fefficient-image-segmentation-using-pytorch-part-4-6c86da083432%2F&title=Efficient%20Image%20Segmentation%20Using%20PyTorch%3A%20Part%204)\n * [](https://www.linkedin.com/shareArticle?mini=true&url=https%3A%2F%2Ftowardsdatascience.com%2Fefficient-image-segmentation-using-pytorch-part-4-6c86da083432%2F&title=Efficient%20Image%20Segmentation%20Using%20PyTorch%3A%20Part%204)\n * [](https://x.com/share?url=https%3A%2F%2Ftowardsdatascience.com%2Fefficient-image-segmentation-using-pytorch-part-4-6c86da083432%2F&text=Efficient%20Image%20Segmentation%20Using%20PyTorch%3A%20Part%204)\n\n\nTowards Data Science is a community publication. Submit your insights to reach our global audience and earn through the TDS Author Payment Program.\n[Write for TDS](https://towardsdatascience.com/submissions/)\n## Related Articles\n * [![](../images/009.jpg)](https://towardsdatascience.com/landing-your-first-machine-learning-job-startup-vs-big-tech-vs-academia/)\n## [Landing your First Machine Learning Job: Startup vs Big Tech vs Academia](https://towardsdatascience.com/landing-your-first-machine-learning-job-startup-vs-big-tech-vs-academia/)\n[Machine Learning](https://towardsdatascience.com/category/artificial-intelligence/machine-learning/)\nA practical guide to landing your first Machine Learning job across startups, big tech, and academia.\n[Piero Paialunga](https://towardsdatascience.com/author/piero-paialunga/)\nJune 3, 202518 min read\n * [![](../images/010.png)](https://towardsdatascience.com/creating-and-training-a-u-net-model-with-pytorch-for-2d-3d-semantic-segmentation-training-3-4-8242d31de234/)\n## [Creating and training a U-Net model with PyTorch for 2D & 3D semantic segmentation: Training [3/4]](https://towardsdatascience.com/creating-and-training-a-u-net-model-with-pytorch-for-2d-3d-semantic-segmentation-training-3-4-8242d31de234/)\n[Deep Learning](https://towardsdatascience.com/category/artificial-intelligence/deep-learning/)\nIn the previous chapters we created our dataset and built the U-Net model. Now it is time to start training. For that we will write our...\n[Johannes Schmidt](https://towardsdatascience.com/author/johschmidt42/)\nDecember 5, 20206 min read\n * [![](../images/011.jpeg)](https://towardsdatascience.com/machine-learning-adventures-with-mlflow-64127713b0a1/)\n## [Machine Learning adventures with MLFlow](https://towardsdatascience.com/machine-learning-adventures-with-mlflow-64127713b0a1/)\n[Machine Learning](https://towardsdatascience.com/category/artificial-intelligence/machine-learning/)\nBuilding ML workflow on your local system\n[Yashaswi Nayak](https://towardsdatascience.com/author/yashaswi_nayak/)\nDecember 20, 202111 min read\n * [![](../images/012.jpg)](https://towardsdatascience.com/ai-model-optimization-on-aws-inferentia-and-trainium-cfd48e85d5ac/)\n## [AI Model Optimization on AWS Inferentia and Trainium](https://towardsdatascience.com/ai-model-optimization-on-aws-inferentia-and-trainium-cfd48e85d5ac/)\n[Artificial Intelligence](https://towardsdatascience.com/category/artificial-intelligence/)\nTips for accelerating ML with AWS Neuron SDK\n[Chaim Rand](https://towardsdatascience.com/author/chaimrand/)\nOctober 20, 202410 min read\n * [![](../images/013.jpg)](https://towardsdatascience.com/what-is-quantum-machine-learning-and-how-to-create-a-quantum-neural-network-model-a59bd571021d/)\n## [How to Classify Encrypted Data Using a Quantum Neural Network Model](https://towardsdatascience.com/what-is-quantum-machine-learning-and-how-to-create-a-quantum-neural-network-model-a59bd571021d/)\n[Machine Learning](https://towardsdatascience.com/category/artificial-intelligence/machine-learning/)\nA beginner's guide to quantum machine learning\n[Saurav Pawar](https://towardsdatascience.com/author/pawarsaurav842/)\nOctober 23, 20226 min read\n * [![](../images/014.jpeg)](https://towardsdatascience.com/transformer-unleashed-deep-forecasting-of-multivariate-time-series-in-python-9ca729dac019/)\n## [Transformer Unleashed: Deep Forecasting of Multivariate Time Series in Python](https://towardsdatascience.com/transformer-unleashed-deep-forecasting-of-multivariate-time-series-in-python-9ca729dac019/)\n[](https://towardsdatascience.com/efficient-image-segmentation-using-pytorch-part-4-6c86da083432/)\nEnd-to-End Example: Probabilistic Forecast of a Multivariate Time Series with Complex Seasonality\n[Heiko Onnen](https://towardsdatascience.com/author/h3ik0.th/)\nJanuary 5, 20226 min read\n * [![](../images/015.png)](https://towardsdatascience.com/causal-inference-example-elasticity-de4a3e2e621b/)\n## [Causal Inference in the Wild](https://towardsdatascience.com/causal-inference-example-elasticity-de4a3e2e621b/)\n[Machine Learning](https://towardsdatascience.com/category/artificial-intelligence/machine-learning/)\nElasticity pricing as a real-world example of machine-learning-powered causal inference\n[Lars Roemheld](https://towardsdatascience.com/author/lars.roemheld/)\nJuly 8, 202114 min read\n * [![](../images/016.jpg)](https://towardsdatascience.com/python-type-hinting-with-literal-03c60ce42750/)\n## [Python Type Hinting with Literal](https://towardsdatascience.com/python-type-hinting-with-literal-03c60ce42750/)\n[Programming](https://towardsdatascience.com/category/programming/)\nMore powerful than it seems: use typing.Literal to create literal types\n[Marcin Kozak](https://towardsdatascience.com/author/nyggus/)\nNovember 28, 202316 min read\n * [![](../images/017.png)](https://towardsdatascience.com/simplifying-transformers-state-of-the-art-nlp-using-words-you-understand-part-5-decoder-and-cd2810c6ad40/)\n## [Simplifying Transformers: State of the Art NLP Using Words You Understand - part 5- Decoder and...](https://towardsdatascience.com/simplifying-transformers-state-of-the-art-nlp-using-words-you-understand-part-5-decoder-and-cd2810c6ad40/)\n[Deep Learning](https://towardsdatascience.com/category/artificial-intelligence/deep-learning/)\nLast part of the Transformer series\n[Chen Margalit](https://towardsdatascience.com/author/chenmargalit/)\nOctober 5, 20237 min read\n\n\n * [YouTube](https://www.youtube.com/c/TowardsDataScience)\n * [X](https://x.com/TDataScience)\n * [LinkedIn](https://www.linkedin.com/company/towards-data-science/)\n * [Threads](https://www.threads.net/@towardsdatascience)\n * [Bluesky](https://bsky.app/profile/towardsdatascience.com)\n\n\n[](https://towardsdatascience.com/)\nYour home for data science and AI. The world's leading publication for data science, data analytics, data engineering, machine learning, and artificial intelligence professionals.\nSubscribe to Our Newsletter[Write For TDS](https://towardsdatascience.com/submissions/)•[About](https://towardsdatascience.com/about-towards-data-science-d691af11cc2f/)•[Advertise](https://contact.towardsdatascience.com/advertise-with-towards-data-science)•[Privacy Policy](https://towardsdatascience.com/privacy-policy/)•[Terms of Use](https://towardsdatascience.com/website-terms-of-use/)\n© Insight Media Group, LLC 2026•[The New Stack](https://thenewstack.io)•[roadmap.sh](https://roadmap.sh)•[Creator Network](https://insightmediagroup.io/creator-network)\n#### Recommended Articles \nClose\n * [ ![](https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/607639/image?etag=1781250125848&width=500&height=500) Javascript Fatigue: HTMX Is All You Need to Build ChatGPT — Part 1 ](https://towardsdatascience.com/javascript-fatigue-you-dont-need-js-to-build-chatgpt/ \"Javascript Fatigue: HTMX Is All You Need to Build ChatGPT — Part 1\")\n * [ ![](https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/607930/image?etag=1782124566054&width=500&height=500) The Subset Sum Problem Solved in Linear Time for Dense Enough Inputs ](https://towardsdatascience.com/subset-sum-problem-solved-in-linear-time-for-dense-enough-inputs/ \"The Subset Sum Problem Solved in Linear Time for Dense Enough Inputs\")\n * [ ![](https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/3660/image?etag=1780176246070&width=500&height=500) Does Your Company Have a Data Strategy? ](https://towardsdatascience.com/does-your-company-have-a-data-strategy-e8944dbec4ae/ \"Does Your Company Have a Data Strategy?\")\n * [ ![](https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/4970/image?etag=1780565225463&width=500&height=500) Productivity Tips, Data Career Insights, and Other Recent Must-Reads ](https://towardsdatascience.com/productivity-tips-data-career-insights-and-other-recent-must-reads-5c47eb7a2fe0/ \"Productivity Tips, Data Career Insights, and Other Recent Must-Reads\")\n\n\n", "metadata": { "title": "Efficient Image Segmentation Using PyTorch: Part 4 | Towards Data Science", "description": "A Vision Transformer-based model", "keywords": null, "author": "Dhruv Matani", "og:locale": "en_US", "og:type": "article", "og:title": "Efficient Image Segmentation Using PyTorch: Part 4", "og:description": "A Vision Transformer-based model", "og:url": "https://towardsdatascience.com/efficient-image-segmentation-using-pytorch-part-4-6c86da083432/", "og:site_name": "Towards Data Science", "og:image": "https://assets.insightmediagroup.io/media/wp-content/uploads/2023/06/0c3rF_3Zpfv6CCUxq.png", "og:image:width": "696", "og:image:height": "405", "og:image:type": "image/png", "twitter:card": "summary_large_image", "twitter:site": "@TDataScience", "twitter:creator": "@TDataScience", "twitter:label1": "Written by", "twitter:data1": "Dhruv Matani", "twitter:label2": "Est. reading time", "twitter:data2": "15 min read", "article:published_time": "2023-06-27T05:13:30.000Z", "article:modified_time": "2023-06-27T05:13:30.000Z", "article:author": "https://towardsdatascience.com/author/dhruvbird/" }, "paywall": { "listed_paid": false, "detected": false, "marker": "", "public_only": false, "policy": "付费/会员内容只保留页面公开可见部分，不绕过付费墙" }, "crawl": { "success": true, "status_code": 200, "error": "" }, "code_blocks": [ { "language": "txt", "file": "attachments/d5fe5f752277/code/block_000.txt", "chars": 310 }, { "language": "txt", "file": "attachments/d5fe5f752277/code/block_001.txt", "chars": 314 } ], "images": [ { "status": "http_404", "url": "https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/607639/image?etag=1781250125848&width=500&height=500", "src": "https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/607639/image?etag=1781250125848&width=500&height=500", "alt": "", "downloaded": false, "discovery": "saved_markdown", "retried_at": "2026-10-20T17:55:56Z" }, { "status": "http_404", "url": "https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/607930/image?etag=1782124566054&width=500&height=500", "src": "https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/607930/image?etag=1782124566054&width=500&height=500", "alt": "", "downloaded": false, "discovery": "saved_markdown", "retried_at": "2026-10-20T17:55:56Z" }, { "status": "http_404", "url": "https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/3660/image?etag=1780176246070&width=500&height=500", "src": "https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/3660/image?etag=1780176246070&width=500&height=500", "alt": "", "downloaded": false, "discovery": "saved_markdown", "retried_at": "2026-10-20T17:55:56Z" }, { "status": "http_404", "url": "https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/4970/image?etag=1780565225463&width=500&height=500", "src": "https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/4970/image?etag=1780565225463&width=500&height=500", "alt": "", "downloaded": false, "discovery": "saved_markdown", "retried_at": "2026-10-20T17:55:56Z" } ], "audio": [], "videos": [ { "url": "https://www.youtube.com/watch?v=kCc8FmEb1nY", "status": "ok", "local_path": "attachments/d5fe5f752277/videos/repair-kCc8FmEb1nY.mp4", "bytes": 372815482, "valid": true, "height": 1080, "width": 1920, "duration": "6979.901000" }, { "url": "https://www.youtube.com/c/TowardsDataScience", "kind": "channel_reference", "status": "reference_only" } ], "links": [ { "type": "internal", "href": "https://towardsdatascience.com/" }, { "type": "internal", "href": "https://towardsdatascience.com/latest/" }, { "type": "internal", "href": "https://towardsdatascience.com/tag/editors-pick/" }, { "type": "internal", "href": "https://towardsdatascience.com/tag/the-variable/" }, { "type": "internal", "href": "https://towardsdatascience.com/shipai/" }, { "type": "internal", "href": "https://towardsdatascience.com/submissions/" }, { "type": "internal", "href": "https://towardsdatascience.com/author/dhruvbird/" }, { "type": "internal", "href": "https://towardsdatascience.com/replace-manual-normalization-with-batch-normalization-in-vision-ai-models-e7782e82193c" }, { "type": "internal", "href": "https://towardsdatascience.com/implementing-visualttransformer-in-pytorch-184f9f16f632" }, { "type": "internal", "href": "https://towardsdatascience.com/implementing-segformer-in-pytorch-8f4705e2ed0e" }, { "type": "internal", "href": "https://towardsdatascience.com/tag/deep-dives/" }, { "type": "internal", "href": "https://towardsdatascience.com/tag/pytorch/" }, { "type": "internal", "href": "https://towardsdatascience.com/tag/pytorch-image-segment/" }, { "type": "internal", "href": "https://towardsdatascience.com/tag/transformers/" }, { "type": "internal", "href": "https://towardsdatascience.com/tag/vision-transformer/" }, { "type": "internal", "href": "https://towardsdatascience.com/landing-your-first-machine-learning-job-startup-vs-big-tech-vs-academia/" }, { "type": "internal", "href": "https://towardsdatascience.com/category/artificial-intelligence/machine-learning/" }, { "type": "internal", "href": "https://towardsdatascience.com/author/piero-paialunga/" }, { "type": "internal", "href": "https://towardsdatascience.com/creating-and-training-a-u-net-model-with-pytorch-for-2d-3d-semantic-segmentation-training-3-4-8242d31de234/" }, { "type": "internal", "href": "https://towardsdatascience.com/category/artificial-intelligence/deep-learning/" }, { "type": "external", "href": "https://write.insightmediagroup.io" }, { "type": "external", "href": "https://www.linkedin.com/company/towards-data-science/" }, { "type": "external", "href": "https://x.com/TDataScience" }, { "type": "external", "href": "https://medium.com/@brocolishbroxoli" }, { "type": "external", "href": "https://ai.googleblog.com/2017/08/transformer-novel-neural-network.html" }, { "type": "external", "href": "https://github.com/dhruvbird/ml-notebooks/blob/main/pets_segmentation/vision-transformer-for-oxford-iiit-pet-segmentatio.ipynb" }, { "type": "external", "href": "https://medium.com/p/89e8297a0923/" }, { "type": "external", "href": "https://medium.com/p/bed68cadd7c7/" }, { "type": "external", "href": "https://medium.com/p/3534cf04fb89/" }, { "type": "external", "href": "https://arxiv.org/abs/1706.03762" }, { "type": "external", "href": "https://pytorch.org/docs/stable/generated/torch.nn.MultiheadAttention.html" }, { "type": "external", "href": "https://colab.research.google.com/github/tensorflow/tensor2tensor/blob/master/tensor2tensor/notebooks/hello_t2t.ipynb" }, { "type": "external", "href": "https://arxiv.org/abs/2010.11929" }, { "type": "external", "href": "https://arxiv.org/pdf/2010.11929.pdf" }, { "type": "external", "href": "https://github.com/dhruvbird/ml-notebooks/blob/main/pets_segmentation/Vision%20Transformer%20Article%20Images.ipynb" }, { "type": "external", "href": "https://en.wikipedia.org/wiki/Multilayer_perceptron" }, { "type": "external", "href": "https://arxiv.org/pdf/2105.15203.pdf" }, { "type": "external", "href": "http://jalammar.github.io/illustrated-transformer/" }, { "type": "external", "href": "https://www.youtube.com/watch?v=kCc8FmEb1nY" }, { "type": "external", "href": "https://einops.rocks/" } ], "repair_version": 2, "repaired_at": "2026-10-20T17:28:33Z", "source_url": "https://towardsdatascience.com/efficient-image-segmentation-using-pytorch-part-4-6c86da083432/", "delivery_status": "unverified", "discovery_status": "http_403", "missing_reason": "源页面未重新核验:http_403；已按保留正文补采图片", "repair_error": "", "actual_video_files": [ { "local_path": "attachments/d5fe5f752277/videos/repair-kCc8FmEb1nY.mp4", "bytes": 372815482, "valid": true, "height": 1080, "width": 1920, "duration": "6979.901000", "status": "ok" } ], "expected_saved_body_images": 4, "image_discovery": "saved_markdown", "source_body_markdown": "![Revisit consent button](https://cdn-cookieyes.com/assets/images/revisit.svg)\nWe value your privacy\nWe use cookies to enhance your browsing experience, serve personalised ads or content, and analyse our traffic. By clicking \"Accept All\", you consent to our use of cookies.\nCustomise Reject All Accept All\nCustomise Consent Preferences ![](https://cdn-cookieyes.com/assets/images/close.svg)\nWe use cookies to help you navigate efficiently and perform certain functions. You will find detailed information about all cookies under each consent category below.\nThe cookies that are categorised as \"Necessary\" are stored on your browser as they are essential for enabling the basic functionalities of the site. ... Show more\nNecessaryAlways Active\nNecessary cookies are required to enable the basic features of this site, such as providing secure log-in or adjusting your consent preferences. These cookies do not store any personally identifiable data.\n * Cookie\nBCTempID\n * Duration\n10 minutes\n * Description\nNo description available.\n\n\n * Cookie\n__cf_bm\n * Duration\n1 hour\n * Description\nThis cookie, set by Cloudflare, is used to support Cloudflare Bot Management. \n\n\n * Cookie\nAWSALBCORS\n * Duration\n7 days\n * Description\nAmazon Web Services set this cookie for load balancing.\n\n\n * Cookie\n_cfuvid\n * Duration\nsession\n * Description\nCloudflare sets this cookie to track users across sessions to optimize user experience by maintaining session consistency and providing personalized services\n\n\n * Cookie\nli_gc\n * Duration\n6 months\n * Description\nLinkedin set this cookie for storing visitor's consent regarding using cookies for non-essential purposes.\n\n\n * Cookie\n__hssrc\n * Duration\nsession\n * Description\nThis cookie is set by Hubspot whenever it changes the session cookie. The __hssrc cookie set to 1 indicates that the user has restarted the browser, and if the cookie does not exist, it is assumed to be a new session.\n\n\n * Cookie\n__hssc\n * Duration\n1 hour\n * Description\nHubSpot sets this cookie to keep track of sessions and to determine if HubSpot should increment the session number and timestamps in the __hstc cookie.\n\n\n * Cookie\nwpEmojiSettingsSupports\n * Duration\nsession\n * Description\nWordPress sets this cookie when a user interacts with emojis on a WordPress site. It helps determine if the user's browser can display emojis properly.\n\n\n * Cookie\nBCSessionID\n * Duration\n1 year 1 month 4 days\n * Description\nBlueconic sets this cookie as a unique identifier for the BlueConic profile.\n\n\n * Cookie\n_octo\n * Duration\n1 year\n * Description\nNo description available.\n\n\n * Cookie\nlogged_in\n * Duration\n1 year\n * Description\nNo description available.\n\n\n * Cookie\n__Secure-YEC\n * Duration\npast\n * Description\nYouTube sets this cookie to stores the user's video player preferences using embedded YouTube video\n\n\n * Cookie\n__eoi\n * Duration\n6 months\n * Description\nDescription is currently not available.\n\n\n * Cookie\nAWSALBTGCORS\n * Duration\n7 days\n * Description\nNo description available.\n\n\n * Cookie\nlogin-status-p\n * Duration\npast\n * Description\nDescription is currently not available.\n\n\n * Cookie\nAWSALBTG\n * Duration\n7 days\n * Description\nNo description available.\n\n\n * Cookie\ncsrf_token\n * Duration\nsession\n * Description\nNo description available.\n\n\n * Cookie\ntoken_v2\n * Duration\n1 day\n * Description\nDescription is currently not available.\n\n\n * Cookie\nD\n * Duration\n1 year\n * Description\nDescription is currently not available.\n\n\n * Cookie\nPHPSESSID\n * Duration\nsession\n * Description\nThis cookie is native to PHP applications. The cookie stores and identifies a user's unique session ID to manage user sessions on the website. The cookie is a session cookie and will be deleted when all the browser windows are closed.\n\n\n * Cookie\nVISITOR_PRIVACY_METADATA\n * Duration\n6 months\n * Description\nYouTube sets this cookie to store the user's cookie consent state for the current domain. \n\n\n * Cookie\ncookietest\n * Duration\nsession\n * Description\nThe cookietest cookie is typically used to determine whether the user's browser accepts cookies, essential for website functionality and user experience.\n\n\n * Cookie\n__Host-airtable-session\n * Duration\n1 year\n * Description\nThis cookie is used to enable us to integrate the services of Airtable.\n\n\n * Cookie\n__Host-airtable-session.sig\n * Duration\n1 year\n * Description\nThis cookie is used to enable us to integrate the services of Airtable.\n\n\n * Cookie\nm\n * Duration\n1 year 1 month 4 days\n * Description\nStripe sets this cookie for fraud prevention purposes. It identifies the device used to access the website, allowing the website to be formatted accordingly.\n\n\n * Cookie\nBIGipServer*\n * Duration\nsession\n * Description\nMarketo sets this cookie to collect information about the user's online activity and build a profile about their interests to provide advertisements relevant to the user. \n\n\n * Cookie\n__cfruid\n * Duration\nsession\n * Description\nCloudflare sets this cookie to identify trusted web traffic.\n\n\n * Cookie\n_GRECAPTCHA\n * Duration\n6 months\n * Description\nGoogle Recaptcha service sets this cookie to identify bots to protect the website against malicious spam attacks.\n\n\n * Cookie\n__Secure-YNID\n * Duration\n6 months\n * Description\nGoogle cookie used to protect user security and prevent fraud, especially during the login process.\n\n\n * Cookie\ncookieyes-consent\n * Duration\n1 year\n * Description\nCookieYes sets this cookie to remember users' consent preferences so that their preferences are respected on subsequent visits to this site. It does not collect or store any personal information about the site visitors.\n\n\nFunctional\nFunctional cookies help perform certain functionalities like sharing the content of the website on social media platforms, collecting feedback, and other third-party features.\n * Cookie\nlidc\n * Duration\n1 day\n * Description\nLinkedIn sets the lidc cookie to facilitate data center selection.\n\n\n * Cookie\nbrw\n * Duration\n1 year\n * Description\nNo description available.\n\n\n * Cookie\nbrwConsent\n * Duration\n5 minutes\n * Description\nDescription is currently not available.\n\n\n * Cookie\nWMF-Uniq\n * Duration\n1 year\n * Description\nDescription is currently not available.\n\n\n * Cookie\nloom_anon_comment\n * Duration\n1 year\n * Description\nNo description available.\n\n\n * Cookie\nloom_referral_video\n * Duration\nsession\n * Description\nDescription is currently not available.\n\n\n * Cookie\nVISITOR_INFO1_LIVE\n * Duration\n6 months\n * Description\nA cookie set by YouTube to measure bandwidth that determines whether the user gets the new or old player interface.\n\n\n * Cookie\nyt-remote-connected-devices\n * Duration\nNever Expires\n * Description\nYouTube sets this cookie to store the user's video preferences using embedded YouTube videos.\n\n\n * Cookie\nytidb::LAST_RESULT_ENTRY_KEY\n * Duration\nNever Expires\n * Description\nThe cookie ytidb::LAST_RESULT_ENTRY_KEY is used by YouTube to store the last search result entry that was clicked by the user. This information is used to improve the user experience by providing more relevant search results in the future.\n\n\n * Cookie\nyt-remote-device-id\n * Duration\nNever Expires\n * Description\nYouTube sets this cookie to store the user's video preferences using embedded YouTube videos.\n\n\n * Cookie\nyt-remote-session-name\n * Duration\nsession\n * Description\nThe yt-remote-session-name cookie is used by YouTube to store the user's video player preferences using embedded YouTube video.\n\n\n * Cookie\nyt-remote-fast-check-period\n * Duration\nsession\n * Description\nThe yt-remote-fast-check-period cookie is used by YouTube to store the user's video player preferences for embedded YouTube videos.\n\n\n * Cookie\nyt-remote-session-app\n * Duration\nsession\n * Description\nThe yt-remote-session-app cookie is used by YouTube to store user preferences and information about the interface of the embedded YouTube video player.\n\n\n * Cookie\nyt-remote-cast-available\n * Duration\nsession\n * Description\nThe yt-remote-cast-available cookie is used to store the user's preferences regarding whether casting is available on their YouTube video player.\n\n\n * Cookie\nyt-remote-cast-installed\n * Duration\nsession\n * Description\nThe yt-remote-cast-installed cookie is used to store the user's video player preferences using embedded YouTube video.\n\n\n * Cookie\ncp_session\n * Duration\n3 months\n * Description\nCodepen sets this cookie for Help systems found in the website.\n\n\n * Cookie\nloid\n * Duration\n1 year 1 month 4 days\n * Description\nThis cookie is set by the Reddit. The cookie enables the sharing of content from the website onto the social media platform.\n\n\nAnalytics\nAnalytical cookies are used to understand how visitors interact with the website. These cookies help provide information on metrics such as the number of visitors, bounce rate, traffic source, etc.\n * Cookie\n__hstc\n * Duration\n6 months\n * Description\nHubspot set this main cookie for tracking visitors. It contains the domain, initial timestamp (first visit), last timestamp (last visit), current timestamp (this visit), and session number (increments for each subsequent session).\n\n\n * Cookie\nhubspotutk\n * Duration\n6 months\n * Description\nHubSpot sets this cookie to keep track of the visitors to the website. This cookie is passed to HubSpot on form submission and used when deduplicating contacts.\n\n\n * Cookie\n_ga\n * Duration\n1 year 1 month 4 days\n * Description\nGoogle Analytics sets this cookie to calculate visitor, session and campaign data and track site usage for the site's analytics report. The cookie stores information anonymously and assigns a randomly generated number to recognise unique visitors.\n\n\n * Cookie\n_ga_*\n * Duration\n1 year 1 month 4 days\n * Description\nGoogle Analytics sets this cookie to store and count page views.\n\n\n * Cookie\n__Host-psifi.analyticsTrace\n * Duration\n6 hours\n * Description\nDescription is currently not available.\n\n\n * Cookie\n__Host-psifi.analyticsTraceV2\n * Duration\n6 hours\n * Description\nDescription is currently not available.\n\n\n * Cookie\n_gh_sess\n * Duration\nsession\n * Description\nGitHub sets this cookie for temporary application and framework state between pages like what step the user is on in a multiple step form.\n\n\n * Cookie\nYSC\n * Duration\nsession\n * Description\nYSC cookie is set by Youtube and is used to track the views of embedded videos on Youtube pages.\n\n\n * Cookie\najs_anonymous_id\n * Duration\n1 year\n * Description\nThis cookie is set by Segment to count the number of people who visit a certain site by tracking if they have visited before.\n\n\n * Cookie\nvuid\n * Duration\n1 year 1 month 4 days\n * Description\nVimeo installs this cookie to collect tracking information by setting a unique ID to embed videos on the website. \n\n\nPerformance\nPerformance cookies are used to understand and analyse the key performance indexes of the website which helps in delivering a better user experience for the visitors.\n * Cookie\nAWSALB\n * Duration\n7 days\n * Description\nAWSALB is an application load balancer cookie set by Amazon Web Services to map the session to the target.\n\n\n * Cookie\nacq\n * Duration\npast\n * Description\nDescription is currently not available.\n\n\n * Cookie\nacq.sig\n * Duration\npast\n * Description\nDescription is currently not available.\n\n\n * Cookie\nptc\n * Duration\n2 years\n * Description\nNo description available.\n\n\nAdvertisement\nAdvertisement cookies are used to provide visitors with customised advertisements based on the pages you visited previously and to analyse the effectiveness of the ad campaigns.\n * Cookie\nmuc_ads\n * Duration\n1 year 1 month 4 days\n * Description\nTwitter sets this cookie to collect user behaviour and interaction data to optimize the website.\n\n\n * Cookie\nguest_id_marketing\n * Duration\n1 year 1 month 4 days\n * Description\nTwitter sets this cookie to identify and track the website visitor.\n\n\n * Cookie\nguest_id_ads\n * Duration\n1 year 1 month 4 days\n * Description\nTwitter sets this cookie to identify and track the website visitor.\n\n\n * Cookie\npersonalization_id\n * Duration\n1 year 1 month 4 days\n * Description\nTwitter sets this cookie to integrate and share features for social media and also store information about how the user uses the website, for tracking and targeting.\n\n\n * Cookie\nguest_id\n * Duration\n1 year 1 month 4 days\n * Description\nTwitter sets this cookie to identify and track the website visitor. It registers if a user is signed in to the Twitter platform and collects information about ad preferences.\n\n\n * Cookie\nbcookie\n * Duration\n1 year\n * Description\nLinkedIn sets this cookie from LinkedIn share buttons and ad tags to recognize browser IDs.\n\n\n * Cookie\n__Secure-ROLLOUT_TOKEN\n * Duration\n6 months\n * Description\nYouTube sets this cookie to manage feature rollout and experimentation. It helps Google control which new features or interface changes are shown to users as part of testing and staged rollouts, ensuring consistent experience for a given user during an experiment.\n\n\n * Cookie\nyt.innertube::nextId\n * Duration\nNever Expires\n * Description\nYouTube sets this cookie to register a unique ID to store data on what videos from YouTube the user has seen.\n\n\n * Cookie\nyt.innertube::requests\n * Duration\nNever Expires\n * Description\nYouTube sets this cookie to register a unique ID to store data on what videos from YouTube the user has seen.\n\n\n * Cookie\nsession_tracker\n * Duration\nsession\n * Description\nThis cookie is set by the Reddit. This cookie is used to identify trusted web traffic. It also helps in adverstising on the website.\n\n\n * Cookie\nedgebucket\n * Duration\nsession\n * Description\nReddit sets this cookie to save the information about a log-on Reddit user, for the purpose of advertisement recommendations and updating the content.\n\n\n * Cookie\ndid\n * Duration\n1 year\n * Description\nArbor sets this cookie to show targeted ads to site visitors.This cookie expires after 2 months or 1 year.\n\n\nUncategorised\nOther uncategorised cookies are those that are being analysed and have not been classified into a category as yet.\nNo cookies to display.\nReject All Save My Preferences Accept All\nToggle Search[](https://towardsdatascience.com/)Toggle Mobile Navigation\n[](https://towardsdatascience.com/)\nPublish AI, ML & data-science insights to a global community of data professionals.\nSign in[Submit an Article](https://write.insightmediagroup.io)\n[Latest](https://towardsdatascience.com/latest/)[Editor's Picks](https://towardsdatascience.com/tag/editors-pick/)[Newsletter](https://towardsdatascience.com/tag/the-variable/)[ShipAI](https://towardsdatascience.com/shipai/)\n* * *\n[Write For TDS](https://towardsdatascience.com/submissions/)\n[LinkedIn](https://www.linkedin.com/company/towards-data-science/)[X](https://x.com/TDataScience)Toggle Search\n# Efficient Image Segmentation Using PyTorch: Part 4\nA Vision Transformer-based model\n[Dhruv Matani](https://towardsdatascience.com/author/dhruvbird/)\nJune 27, 202315 min read\nIn this 4-part series, we'll implement image segmentation step by step from scratch using deep learning techniques in PyTorch. This part will focus on implementing a Vision Transformer based model for image segmentation.\nCo-authored with [Naresh Singh](https://medium.com/@brocolishbroxoli)\n![Figure 1: Result of running image segmentation using a vision transformer model architecture. From top to bottom, input images, ground truth segmentation masks, and predicted segmentation masks. Source: Author\\(s\\)](https://assets.insightmediagroup.io/media/wp-content/uploads/2023/06/0c3rF_3Zpfv6CCUxq.png)\nFigure 1: Result of running image segmentation using a vision transformer model architecture. From top to bottom, input images, ground truth segmentation masks, and predicted segmentation masks. Source: Author(s)\n* * *\n## Article outline\nIn this article, we will visit the [Transformer architecture](https://ai.googleblog.com/2017/08/transformer-novel-neural-network.html) which has taken the world of deep learning by storm. The transformer is a multimodal architecture which can model different modalities such as language, vision, and audio.\nIn this article, we will\n 1. Learn about the transformer architecture and the key concepts involved\n 2. Understand the vision transformer architecture\n 3. Introduce a vision transformer model that is written from scratch so that you can appreciate all the building blocks and moving parts\n 4. Follow an input tensor fed into this model and inspect how it changes shape\n 5. Use this model to perform image segmentation on the Oxford IIIT Pet dataset\n 6. Observe the results of this segmentation task\n 7. Briefly introduce the SegFormer, a state of the art vision transformer for semantic segmentation\n\n\nThroughout this article, we will reference code and results from this [notebook](https://github.com/dhruvbird/ml-notebooks/blob/main/pets_segmentation/vision-transformer-for-oxford-iiit-pet-segmentatio.ipynb) for model training. If you wish to reproduce the results, you'll need a GPU to ensure that the first notebook completes running in a reasonable amount of time.\n* * *\n## Articles in this series\nThis series is for readers at all experience levels with deep learning. If you want to learn about the practice of deep learning and vision AI along with some solid theory and hands-on experience, you've come to the right place! This is expected to be a 4-part series with the following articles:\n 1. [Concepts and Ideas](https://medium.com/p/89e8297a0923/)\n 2. [A CNN-based model](https://medium.com/p/bed68cadd7c7/)\n 3. [Depthwise separable convolutions](https://medium.com/p/3534cf04fb89/)\n 4. **A Vision Transformer-based model (this article)**\n\n\nLet's start our journey into vision transformers with an introduction and intuitive understanding of the transformer architecture.\n* * *\n## The Transformer Architecture\nWe can think of the transformer architecture as a composition of interleaving layers of _communication_ and _computation_. This idea is depicted visually in figure 2. The transformer has N processing units (N is 3 in Figure 2), each of which is responsible for processing a 1/N fraction of the input. For those processing units to produce meaningful results, each of them need to have a _global view_ of the input. Hence, the system repeatedly communicates information about the data in every processing unit to every other processing unit; shown using the red, green, and blue arrows going from every processing unit to every other processing unit. This is followed by some computation based on this information. After sufficient repetitions of this process, the model is able to produce desired results.\n![Figure 2: Interleaved communication and computation in transformers. The image shows just 2 layers of communication and computation. In practice, there are many more such layers. Source: Author\\(s\\).](https://assets.insightmediagroup.io/media/wp-content/uploads/2023/06/0bpWZaQdbzM4lT7dx.png)\nFigure 2: Interleaved communication and computation in transformers. The image shows just 2 layers of communication and computation. In practice, there are many more such layers. Source: Author(s).\nIt's worth noting most online resources typically discuss both the encoder and the decoder of the transformer as presented in the paper titled \"[Attention is all you need](https://arxiv.org/abs/1706.03762).\" However, in this article, we will describe just the encoder part of the transformer.\nLet's take a closer look at what constitutes communication and computation in transformers.\n### Communication in transformers: Attention\nIn transformers, communication is implemented by a layer known as the attention layer. In PyTorch, this is called [MultiHeadAttention](https://pytorch.org/docs/stable/generated/torch.nn.MultiheadAttention.html). We'll get to the reason for that name in a bit.\nThe documentation says:\n_\"Allows the model to jointly attend to information from different representation subspaces as described in the paper:_[_Attention is all you need_](https://arxiv.org/abs/1706.03762) _.\"_\nThe attention mechanism consumes an input tensor _x_ of shape (Batch, Length, Features), and it produces a similarly shaped tensor _y_ such that the features for each input are updated based on which other inputs in the same instance the tensor is paying attention to. Hence, the features of each tensor of length \"Features\" in the instance of size \"Length\" are updated based on every other tensor. This is where the quadratic cost of the attention mechanism comes in.\n![Figure 3: Attention of the word "it " shown relative to the other words in the sentence. We can see that "it " is paying attention to the words "animal ", "too ", and "tire\\(d\\) " in the same sentence. Source: Generated using this colab.](https://assets.insightmediagroup.io/media/wp-content/uploads/2023/06/09cWryzyUlC6x9m5x.png)\nFigure 3: Attention of the word \"it \" shown relative to the other words in the sentence. We can see that \"it \" is paying attention to the words \"animal \", \"too \", and \"tire(d) \" in the same sentence. Source: Generated using [this colab](https://colab.research.google.com/github/tensorflow/tensor2tensor/blob/master/tensor2tensor/notebooks/hello_t2t.ipynb).\nIn the context of a vision transformer, the input to the transformer is an image. Let's assume this to be a 128 x 128 (width, height) image. We chunk it into multiple smaller patches of size (16 x 16). For a 128 x 128 image, we get 64 patches (Length), 8 patches in each row, and 8 rows of patches.\nEach one of these 64 patches of size 16 x 16 pixels is considered to be a separate input to the transformer model. Without getting too deep into the details, it should be sufficient to think of this process as being driven by 64 different processing units, each of which is processing a single 16x16 image patch.\nIn each round, the attention mechanism in each processing unit is responsible for looking at the image patch it is responsible for and querying each one of the other remaining 63 processing units to ask them for any information that may be relevant and useful to help it effectively process its own image patch.\nThe communication step via attention is followed by computation, which we will look at next.\n### Computation in transformers: Multi Layer Perceptron\nComputation in transformers is nothing but a MultiLayerPerceptron (MLP) unit. This unit is composed of 2 Linear layers, with a GeLU non-linearity in between. One can consider using other non-linearities as well. This unit first projects the input to 4x the size and reprojects it back to 1x, which is the same as the input size.\nIn the code we'll see in our notebook, this class is called MultiLayerPerceptron. The code is shown below.\npython\n\n```\nclass MultiLayerPerceptron(nn.Sequential):\n def __init__(self, embed_size, dropout):\n super().__init__(\n nn.Linear(embed_size, embed_size * 4),\n nn.GELU(),\n nn.Linear(embed_size * 4, embed_size),\n nn.Dropout(p=dropout),\n )\n # end def\n# end class\n```\n\nNow that we understand the high level working of the transformer architecture, let's focus our attention on the vision transformer since we're going to be performing image segmentation.\n* * *\n## The Vision Transformer\nThe vision transformer was first introduced by the paper titled \"[An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://arxiv.org/abs/2010.11929)\". The paper discusses how the authors apply the vanilla transformer architecture to the problem of image classification. This is done by splitting the image into patches of size 16x16, and treating each patch as an input token to the model. The transformer encoder model is fed these input tokens, and is asked to predict a class for the input image.\n![Figure 4: Source: Transformers for image recognition at scale.](https://assets.insightmediagroup.io/media/wp-content/uploads/2023/06/0cuLnz_JK5Afenyrf.png)\nFigure 4: Source: [Transformers for image recognition at scale](https://arxiv.org/pdf/2010.11929.pdf).\nIn our case, we are interested in image segmentation. We can consider it to be a pixel-level classification task because we intend to predict a target class per pixel..\nWe make a small but important change to the vanilla vision transformer and replace the MLP head for classification by an MLP head for pixel level classification. We have a single linear layer in the output that is shared by every patch whose segmentation mask is predicted by the vision transformer. This shared linear layer predicts a segmentation mask for every patch that was sent as input to the model.\nIn the case of the vision transformer, a patch of size 16x16 is considered to be equivalent to a single input token at a specific time step.\n![Figure 5: The end to end working of the vision transformer for image segmentation. Image generated using this notebook. Source: Author\\(s\\).](https://assets.insightmediagroup.io/media/wp-content/uploads/2023/06/0JbCsJbSNF9j_wPVF.png)\nFigure 5: The end to end working of the vision transformer for image segmentation. Image generated using this [notebook](https://github.com/dhruvbird/ml-notebooks/blob/main/pets_segmentation/Vision%20Transformer%20Article%20Images.ipynb). Source: Author(s).\n### Building an intuition for tensor dimensions in vision transformers\nWhen working with deep CNNs, the tensor dimensions we used for the most part was (N, C H, W), where the letters stand for the following:\n * N: Batch size\n * C: Number of channels\n * H: Height\n * W: Width\n\n\nYou can see that this format is geared toward 2d image processing, since it smells of features that are very specific to images.\nWith transformers on the other hand, things become a lot more generic and domain agnostic. What we'll see below applies to vision, text, NLP, audio or other problems where input data can be represented as a sequence. It is worth noting that there's little vision specific bias in the representation of tensors as they flow through our vision transformer.\nWhen working with transformers and attention in general, we expect the tensors to have the following shape: (B, T, C), where the letters stand for the following:\n * B: Batch size (same as that for CNNs)\n * T: Time dimension or sequence length. This dimension is also sometimes called L. In the case of vision transformers, each image patch corresponds to this dimension. If we have 16 image patches, then the value of the T dimension will be 16\n * C: The channel or embedding size dimension. This dimension is also sometimes called E. When processing images, each patch of size 3x16x16 (Channel, Width, Height) is mapped via a patch embedding layer to an embedding of size C. We'll see how this is done later.\n\n\nLet's dive into how the input image tensor gets mutated and processed along its way to predicting the segmentation mask.\n### The journey of a tensor in a vision transformer\nIn deep CNNs, the journey of a tensor looks something like this (in a UNet, SegNet, or other CNN based architecture).\nThe input tensor is typically of shape (1, 3, 128, 128). This tensor goes through a series of convolution and max-pooling operations where its spatial dimensions are reduced and channel dimensions are increased, typically by a factor of 2 each. This is called the feature encoder. After this, we do the reverse operation where we increase the spatial dimensions and reduce the channel dimensions. This is called the feature decoder. After the decoding process, we get a tensor of shape (1, 64, 128, 128). This is then projected into the number of output channels C that we desire as (1, C, 128, 128) using a 1x1 pointwise convolution without bias.\n![Figure 6: Typical progression of tensor shapes through a deep CNN used for image segmentation. Source: Author\\(s\\).](https://assets.insightmediagroup.io/media/wp-content/uploads/2023/06/001xjpWZkVkBlfcjT.png)\nFigure 6: Typical progression of tensor shapes through a deep CNN used for image segmentation. Source: Author(s).\nWith vision transformers, the flow is much more complex. Let's take a look at an image below and then try to understand how the tensor transforms shapes at every step along the way.\n![Figure 7: Typical progression of tensor shapes through a vision transformer for image segmentation. Source: Author\\(s\\).](https://assets.insightmediagroup.io/media/wp-content/uploads/2023/06/0ruzy639FhSDbfNCQ.png)\nFigure 7: Typical progression of tensor shapes through a vision transformer for image segmentation. Source: Author(s).\nLet's look at each step in more detail and see how it updates the shape of the tensor flowing through the vision transformer. To understand this better, let's take concrete values for our tensor dimensions.\n 1. **Batch Normalization:** The input and output tensors have shape (1, 3, 128, 128). The shape is unchanged, but the values are normalized to zero mean and unit variance.\n 2. **Image to patches:** The input tensor of shape (1, 3, 128, 128) is converted into a stacked patch of 16x16 images. The output tensor has shape (1, 64, 768).\n 3. **Patch embedding:** The patch embedding layer maps the 768 input channels to 512 embedding channels (for this example). The output tensor is of shape (1, 64, 512). The patch embedding layer is basically just an nn.Linear layer in PyTorch.\n 4. **Position embedding:** The position embedding layer doesn't have an input tensor, but effectively contributes a learnable parameter (trainable tensor in PyTorch) o f the same shape as the patch embedding. This is of shape (1, 64, 512).\n 5. **Add:** The patch and position embeddings are added together piecewise to produce the input to our vision transformer encoder. This tensor is of shape (1, 64, 512). You'll notice that the main workhorse of the vision transformer, i.e. the encoder basically leaves this tensor shape unchanged.\n 6. **Transformer encoder:** The input tensor of shape (1, 64, 512) flows through multiple transformer encoder blocks, each of which have multiple attention heads (communication) followed by an MLP layer (computation). The tensor shape remains unchanged as (1, 64, 512).\n 7. **Linear output projection:** If we assume that we want to segment each image into 10 classes, then we will need each patch of size 16x16 to have 10 channels. The nn.Linear layer for output projection will now convert the 512 embedding channels to 16x16x10 = 2560 output channels, and this tensor will look like (1, 64, 2560). In the diagram above C' = 10. Ideally, this would be a multi-layer perceptron, since _\"[MLPs are universal function approximators](https://en.wikipedia.org/wiki/Multilayer_perceptron)\"_, but we use a single linear layer since this is an educational exercise\n 8. **Patch to image:** This layer converts the 64 patches encoded as a (1, 64, 2560) tensor back into something that looks like a segmentation mask. This can be 10 single channel images, or in this case a single 10 channel image, with each channel being the segmentation mask for one of the 10 classes. The output tensor is of shape (1, 10, 128, 128).\n\n\nThat's it - we've successfully segmented an input image using a vision transformer! Next, let's take a look at an experiment along with some results.\n### Vision transformers in action\nThis [notebook](https://github.com/dhruvbird/ml-notebooks/blob/main/pets_segmentation/vision-transformer-for-oxford-iiit-pet-segmentatio.ipynb) contains all the code for this section.\nAs far as the code and class structure is concerned, it closely mimics the block diagram above. Most of the concepts mentioned above have a 1:1 correspondence to class names in this [notebook](https://github.com/dhruvbird/ml-notebooks/blob/main/pets_segmentation/vision-transformer-for-oxford-iiit-pet-segmentatio.ipynb).\nThere are some concepts related to the attention layers that are critical hyperparameters for our model. We didn't mention anything about the details of the multi-head attention earlier since we mentioned that it's out of scope for the purposes of this article. We highly recommend reading the reference material mentioned above before proceeding if you don't have a basic understanding of the attention mechanism in transformers.\nWe used the following model parameters for the vision transformer for segmentation.\n 1. 768 embedding dimensions for the PatchEmbedding layer\n 2. 12 Transformer encoder blocks\n 3. 8 attention heads in each transformer encoder block\n 4. 20% dropout in multi-head attention and MLP\n\n\nThis configuration can be seen in the VisionTransformerArgs Python dataclass.\npython\n\n```\n@dataclass\nclass VisionTransformerArgs:\n \"\"\"Arguments to the VisionTransformerForSegmentation.\"\"\"\n image_size: int = 128\n patch_size: int = 16\n in_channels: int = 3\n out_channels: int = 3\n embed_size: int = 768\n num_blocks: int = 12\n num_heads: int = 8\n dropout: float = 0.2\n# end class\n```\n\nA similar configuration as [before](https://medium.com/p/bed68cadd7c7/) was used during model training and validation. The configuration is specified below.\n 1. The _random horizontal flip_ and _colour jitter_ data augmentations are applied to the training set to prevent overfitting\n 2. The images are resized to 128x128 pixels in a non-aspect preserving resize operation\n 3. No input normalization is applied to the images - instead a [batch normalization layer is used as the first layer of the model](https://towardsdatascience.com/replace-manual-normalization-with-batch-normalization-in-vision-ai-models-e7782e82193c)\n 4. The model is trained for 50 epochs using the Adam optimizer with a LR of 0.0004 and a StepLR scheduler that decays the learning rate by 0.8x every 12 epochs\n 5. The cross-entropy loss function is used to classify a pixel as belonging to a pet, the background, or a pet border\n\n\nThe model has 86.28M parameters and achieved a validation accuracy of 85.89% after 50 training epochs. This is less than the 88.28% accuracy achieved by deep CNN model after 20 training epochs. This could be due to a few factors that need to be validated experimentally.\n 1. The last output projection layer is a single nn.Linear and not a multi-layer perceptron\n 2. The 16x16 patch size is too large to capture more fine grained detail\n 3. Not enough training epochs\n 4. Not enough training data - it's known that transformer models need a lot more data to train effectively compared to deep CNN models\n 5. The learning rate is too low\n\n\nWe plotted a gif showing how the model is learning to predict the segmentation masks for 21 images in the validation set.\n![Figure 8: A gif showing the progression of segmentation masks predicted by the vision transformer for image segmentation model. Source: Author\\(s\\).](https://assets.insightmediagroup.io/media/0*lcQNZVJ7F3UVqWuz)\nFigure 8: A gif showing the progression of segmentation masks predicted by the vision transformer for image segmentation model. Source: Author(s).\nWe notice something interesting in the early training epochs. The predicted segmentation masks have some strange blocking artifacts. The only reason we could think of for this is because we're breaking down the image into patches of size 16x16 and after very few training epochs, the model hasn't learned anything useful beyond some very coarse grained information regarding whether this 16x16 patch is generally covered by a pet or by background pixels.\n![Figure 9: The blocking artifacts seen in the predicted segmentation masks when using the vision transformer for image segmentation. Source: Author\\(s\\).](https://assets.insightmediagroup.io/media/wp-content/uploads/2023/06/0DxR-tXo4kmKxx9ND.png)\nFigure 9: The blocking artifacts seen in the predicted segmentation masks when using the vision transformer for image segmentation. Source: Author(s).\nNow that we have seen a basic vision transformer in action, let's turn our attention to a state of the art vision transformer for segmentation tasks.\n### SegFormer: Semantic segmentation with transformers\nThe SegFormer architecture was proposed in [this paper](https://arxiv.org/pdf/2105.15203.pdf) in 2021. The transformer we saw above is a simpler version of the SegFormer architecture.\n![Figure 10: The SegFormer architecture. Source: SegFormer paper \\(2021\\).](https://assets.insightmediagroup.io/media/wp-content/uploads/2023/06/0Xt4yS8unp-aRi9so.png)\nFigure 10: The SegFormer architecture. Source: [SegFormer paper (2021)](https://arxiv.org/pdf/2105.15203.pdf).\nMost notably, the SegFormer:\n 1. Generates 4 sets of images with patches of size 4x4, 8x8, 16x16, and 32x32 instead of a single patched image with patches of size 16x16\n 2. Uses 4 transformer encoder blocks instead of just 1. This feels like a model ensemble\n 3. Uses convolutions in the pre and post phases of self-attention\n 4. Doesn't use positional embeddings\n 5. Each transformer block processes images at spatial resolution H/4 x W/4, H/8 x W/8, H/16 x W/16, and H/32, W/32\n 6. Similarly, the channels increase when the spatial dimensions reduce. This feels similar to deep CNNs\n 7. Predictions at multiple spatial dimensions are upsampled and then merged together in the decoder\n 8. An MLP combines all these predictions to provide a final prediction\n 9. The final prediction is at spatial dimension H/4, W/4 and not at H, W.\n\n\n* * *\n## Conclusion\nIn part-4 of this series, we were introduced to the transformer architecture and vision transformers in particular. We developed an intuitive understanding of how vision transformers work, and the basic building block involved in the communication and computation phases of vision transformers. We saw the unique patch based approach adopted by vision transformers for predicting segmentation masks and then combining the predictions together.\nWe reviewed an experiment that shows vision transformers in action, and were able to compare results with deep CNN approaches. While our vision transformer is not state of the art, it was able to achieve pretty decent results. We provided a glimpse into state-of-the-art approaches such as SegFormer.\nIt should be clear by now that transformers have a lot more moving parts and are more complex compared to deep CNN-based approaches. From a raw FLOPs point of view, transformers hold the promise of being more efficient. In transformers, the only real layer that is computationally heavy is nn.Linear. This is implemented using optimized matrix multiplication on most architectures. Due to this architectural simplicity, transformers hold the promise of being easier to optimize and speed up compared to deep CNN-based approaches.\nCongratulations on making it this far! We're glad you enjoyed reading this series on efficient image segmentation in PyTorch. If you have questions or comments, please feel free to leave them in the comments section.\n* * *\n## Further reading\nThe details of the attention mechanism are out of scope for this article. Besides, there are numerous high quality resources that you can refer to to understand the attention mechanism in great detail. Here are some that we highly recommend.\n 1. [The Illustrated Transformer](http://jalammar.github.io/illustrated-transformer/)\n 2. [NanoGPT from scratch using PyTorch](https://www.youtube.com/watch?v=kCc8FmEb1nY)\n\n\nWe'll provide links to articles that provide more details on vision transformers below.\n 1. [Implementing Vision Transformer (ViT) in PyTorch](https://towardsdatascience.com/implementing-visualttransformer-in-pytorch-184f9f16f632): This article details the implementation of a vision transformer for image classification in PyTorch. Notably, their implementation uses [einops](https://einops.rocks/), which we avoid, since this is an education-focused exercise (we recommend learning and using einops for code readability though). We instead use native PyTorch operators for permuting and rearranging tensor dimensions. Additionally, there are a few places where the author uses Conv2d instead of Linear layers. We wanted to build an implementation of vision transformers without the use of convolutional layers entirely.\n 2. [Vision Transformer: AI Summer](https://theaisummer.com/vision-transformer/)\n 3. [Implementing SegFormer in PyTorch](https://towardsdatascience.com/implementing-segformer-in-pytorch-8f4705e2ed0e)\n\n\n## Register Now to Continue Reading\n### Unlock Towards Data Science with a Free Account\nAll articles on Towards Data Science are available to registered users. Create a free account today to access this article and thousands more.\nRegister for freeSign in to your account\n* * *\n![](https://assets.insightmediagroup.io/media/1787657076238_edr3rv.jpeg)\nWritten By\n## [Dhruv Matani](https://towardsdatascience.com/author/dhruvbird/)\n[See all from Dhruv Matani](https://towardsdatascience.com/author/dhruvbird/)\n[Deep Dives](https://towardsdatascience.com/tag/deep-dives/)[Pytorch](https://towardsdatascience.com/tag/pytorch/)[Pytorch Image Segment](https://towardsdatascience.com/tag/pytorch-image-segment/)[Transformers](https://towardsdatascience.com/tag/transformers/)[Vision Transformer](https://towardsdatascience.com/tag/vision-transformer/)\nShare This Article\n * [](https://www.facebook.com/sharer/sharer.php?u=https%3A%2F%2Ftowardsdatascience.com%2Fefficient-image-segmentation-using-pytorch-part-4-6c86da083432%2F&title=Efficient%20Image%20Segmentation%20Using%20PyTorch%3A%20Part%204)\n * [](https://www.linkedin.com/shareArticle?mini=true&url=https%3A%2F%2Ftowardsdatascience.com%2Fefficient-image-segmentation-using-pytorch-part-4-6c86da083432%2F&title=Efficient%20Image%20Segmentation%20Using%20PyTorch%3A%20Part%204)\n * [](https://x.com/share?url=https%3A%2F%2Ftowardsdatascience.com%2Fefficient-image-segmentation-using-pytorch-part-4-6c86da083432%2F&text=Efficient%20Image%20Segmentation%20Using%20PyTorch%3A%20Part%204)\n\n\nTowards Data Science is a community publication. Submit your insights to reach our global audience and earn through the TDS Author Payment Program.\n[Write for TDS](https://towardsdatascience.com/submissions/)\n## Related Articles\n * [![](https://assets.insightmediagroup.io/media/wp-content/uploads/2025/06/heading_pic-scaled-1.jpg)](https://towardsdatascience.com/landing-your-first-machine-learning-job-startup-vs-big-tech-vs-academia/)\n## [Landing your First Machine Learning Job: Startup vs Big Tech vs Academia](https://towardsdatascience.com/landing-your-first-machine-learning-job-startup-vs-big-tech-vs-academia/)\n[Machine Learning](https://towardsdatascience.com/category/artificial-intelligence/machine-learning/)\nA practical guide to landing your first Machine Learning job across startups, big tech, and academia.\n[Piero Paialunga](https://towardsdatascience.com/author/piero-paialunga/)\nJune 3, 202518 min read\n * [![](https://assets.insightmediagroup.io/media/wp-content/uploads/2020/12/1-3w_mXuedRBnSdotpzQ-UQ.png)](https://towardsdatascience.com/creating-and-training-a-u-net-model-with-pytorch-for-2d-3d-semantic-segmentation-training-3-4-8242d31de234/)\n## [Creating and training a U-Net model with PyTorch for 2D & 3D semantic segmentation: Training [3/4]](https://towardsdatascience.com/creating-and-training-a-u-net-model-with-pytorch-for-2d-3d-semantic-segmentation-training-3-4-8242d31de234/)\n[Deep Learning](https://towardsdatascience.com/category/artificial-intelligence/deep-learning/)\nIn the previous chapters we created our dataset and built the U-Net model. Now it is time to start training. For that we will write our...\n[Johannes Schmidt](https://towardsdatascience.com/author/johschmidt42/)\nDecember 5, 20206 min read\n * [![](https://assets.insightmediagroup.io/media/wp-content/uploads/2021/12/1Ia19RmwXhtie6Atab3fODg.jpeg)](https://towardsdatascience.com/machine-learning-adventures-with-mlflow-64127713b0a1/)\n## [Machine Learning adventures with MLFlow](https://towardsdatascience.com/machine-learning-adventures-with-mlflow-64127713b0a1/)\n[Machine Learning](https://towardsdatascience.com/category/artificial-intelligence/machine-learning/)\nBuilding ML workflow on your local system\n[Yashaswi Nayak](https://towardsdatascience.com/author/yashaswi_nayak/)\nDecember 20, 202111 min read\n * [![](https://assets.insightmediagroup.io/media/wp-content/uploads/2024/10/02Rv2SBFg0AgINmyy-scaled.jpg)](https://towardsdatascience.com/ai-model-optimization-on-aws-inferentia-and-trainium-cfd48e85d5ac/)\n## [AI Model Optimization on AWS Inferentia and Trainium](https://towardsdatascience.com/ai-model-optimization-on-aws-inferentia-and-trainium-cfd48e85d5ac/)\n[Artificial Intelligence](https://towardsdatascience.com/category/artificial-intelligence/)\nTips for accelerating ML with AWS Neuron SDK\n[Chaim Rand](https://towardsdatascience.com/author/chaimrand/)\nOctober 20, 202410 min read\n * [![](https://assets.insightmediagroup.io/media/wp-content/uploads/2022/10/02pJ0WNMH_P-xVDsm.jpg)](https://towardsdatascience.com/what-is-quantum-machine-learning-and-how-to-create-a-quantum-neural-network-model-a59bd571021d/)\n## [How to Classify Encrypted Data Using a Quantum Neural Network Model](https://towardsdatascience.com/what-is-quantum-machine-learning-and-how-to-create-a-quantum-neural-network-model-a59bd571021d/)\n[Machine Learning](https://towardsdatascience.com/category/artificial-intelligence/machine-learning/)\nA beginner's guide to quantum machine learning\n[Saurav Pawar](https://towardsdatascience.com/author/pawarsaurav842/)\nOctober 23, 20226 min read\n * [![](https://assets.insightmediagroup.io/media/wp-content/uploads/2022/01/1ZJynP4ii_gTYfvPFzHsogg.jpeg)](https://towardsdatascience.com/transformer-unleashed-deep-forecasting-of-multivariate-time-series-in-python-9ca729dac019/)\n## [Transformer Unleashed: Deep Forecasting of Multivariate Time Series in Python](https://towardsdatascience.com/transformer-unleashed-deep-forecasting-of-multivariate-time-series-in-python-9ca729dac019/)\n[](https://towardsdatascience.com/efficient-image-segmentation-using-pytorch-part-4-6c86da083432/)\nEnd-to-End Example: Probabilistic Forecast of a Multivariate Time Series with Complex Seasonality\n[Heiko Onnen](https://towardsdatascience.com/author/h3ik0.th/)\nJanuary 5, 20226 min read\n * [![](https://assets.insightmediagroup.io/media/wp-content/uploads/2021/07/1oU_kBTmBDCkxanew1vVn0g.png)](https://towardsdatascience.com/causal-inference-example-elasticity-de4a3e2e621b/)\n## [Causal Inference in the Wild](https://towardsdatascience.com/causal-inference-example-elasticity-de4a3e2e621b/)\n[Machine Learning](https://towardsdatascience.com/category/artificial-intelligence/machine-learning/)\nElasticity pricing as a real-world example of machine-learning-powered causal inference\n[Lars Roemheld](https://towardsdatascience.com/author/lars.roemheld/)\nJuly 8, 202114 min read\n * [![](https://assets.insightmediagroup.io/media/wp-content/uploads/2023/11/0hMdXAA2WHIuewfXu-scaled.jpg)](https://towardsdatascience.com/python-type-hinting-with-literal-03c60ce42750/)\n## [Python Type Hinting with Literal](https://towardsdatascience.com/python-type-hinting-with-literal-03c60ce42750/)\n[Programming](https://towardsdatascience.com/category/programming/)\nMore powerful than it seems: use typing.Literal to create literal types\n[Marcin Kozak](https://towardsdatascience.com/author/nyggus/)\nNovember 28, 202316 min read\n * [![](https://assets.insightmediagroup.io/media/wp-content/uploads/2023/10/1cgfP6rpOhtOmxCgricoo7A.png)](https://towardsdatascience.com/simplifying-transformers-state-of-the-art-nlp-using-words-you-understand-part-5-decoder-and-cd2810c6ad40/)\n## [Simplifying Transformers: State of the Art NLP Using Words You Understand - part 5- Decoder and...](https://towardsdatascience.com/simplifying-transformers-state-of-the-art-nlp-using-words-you-understand-part-5-decoder-and-cd2810c6ad40/)\n[Deep Learning](https://towardsdatascience.com/category/artificial-intelligence/deep-learning/)\nLast part of the Transformer series\n[Chen Margalit](https://towardsdatascience.com/author/chenmargalit/)\nOctober 5, 20237 min read\n\n\n * [YouTube](https://www.youtube.com/c/TowardsDataScience)\n * [X](https://x.com/TDataScience)\n * [LinkedIn](https://www.linkedin.com/company/towards-data-science/)\n * [Threads](https://www.threads.net/@towardsdatascience)\n * [Bluesky](https://bsky.app/profile/towardsdatascience.com)\n\n\n[](https://towardsdatascience.com/)\nYour home for data science and AI. The world's leading publication for data science, data analytics, data engineering, machine learning, and artificial intelligence professionals.\nSubscribe to Our Newsletter[Write For TDS](https://towardsdatascience.com/submissions/)•[About](https://towardsdatascience.com/about-towards-data-science-d691af11cc2f/)•[Advertise](https://contact.towardsdatascience.com/advertise-with-towards-data-science)•[Privacy Policy](https://towardsdatascience.com/privacy-policy/)•[Terms of Use](https://towardsdatascience.com/website-terms-of-use/)\n© Insight Media Group, LLC 2026•[The New Stack](https://thenewstack.io)•[roadmap.sh](https://roadmap.sh)•[Creator Network](https://insightmediagroup.io/creator-network)\n#### Recommended Articles \nClose\n * [ ![](https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/607639/image?etag=1781250125848&width=500&height=500) Javascript Fatigue: HTMX Is All You Need to Build ChatGPT — Part 1 ](https://towardsdatascience.com/javascript-fatigue-you-dont-need-js-to-build-chatgpt/ \"Javascript Fatigue: HTMX Is All You Need to Build ChatGPT — Part 1\")\n * [ ![](https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/607930/image?etag=1782124566054&width=500&height=500) The Subset Sum Problem Solved in Linear Time for Dense Enough Inputs ](https://towardsdatascience.com/subset-sum-problem-solved-in-linear-time-for-dense-enough-inputs/ \"The Subset Sum Problem Solved in Linear Time for Dense Enough Inputs\")\n * [ ![](https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/3660/image?etag=1780176246070&width=500&height=500) Does Your Company Have a Data Strategy? ](https://towardsdatascience.com/does-your-company-have-a-data-strategy-e8944dbec4ae/ \"Does Your Company Have a Data Strategy?\")\n * [ ![](https://h030.towardsdatascience.com/rest/v2/contentStores/d7d3a4ab-7caf-4383-bbd2-fbf7fb4f1787/items/4970/image?etag=1780565225463&width=500&height=500) Productivity Tips, Data Career Insights, and Other Recent Must-Reads ](https://towardsdatascience.com/productivity-tips-data-career-insights-and-other-recent-must-reads-5c47eb7a2fe0/ \"Productivity Tips, Data Career Insights, and Other Recent Must-Reads\")\n\n\n", "body_base_path": "attachments/d5fe5f752277/pages/", "body_markdown_path": "attachments/d5fe5f752277/pages/archived-body.md", "local_html": "attachments/d5fe5f752277/pages/archived-body.html", "body_source": "prior_capture" }