酷酷是懒虫

这个屌丝很懒，什么也没留下！

热门标签

Vision Transformer详解（附代码）

作者：酷酷是懒虫 | 2024-06-28 00:17:57

踩

vision transformer

1 引言

$\mathrm{Transformer}$ 在 $\mathrm{NLP}$ 中大获成功， $\mathrm{Vision\text{ }Transformer}$ 则将 $\mathrm{Transformer}$ 模型架构扩展到计算机视觉的领域中，并且它可以很好的地取代卷积操作，在不依赖卷积的情况下，依然可以在图像分类任务上达到很好的效果。卷积操作只能考虑到局部的特征信息，而 $\mathrm{Transformer}$ 中的注意力机制可以综合考量全局的特征信息。 $\mathrm{Vision\text{ }Transformer}$ 尽力做到在不改变 $\mathrm{Transformer}$ 中 $\mathrm{Encoder}$ 架构的前提下，直接将其从 $\mathrm{NLP}$ 领域迁移到计算机视觉领域中，目的是让原始的 $\mathrm{Transformer}$ 模型开箱即用。如果想要了解 $\mathrm{Transformer}$ 原理详细的介绍可以看我的上一篇文章《Transformer详解（附代码）》。

2 注意力机制应用

在正式详细介绍 $\mathrm{Vision\text{ }Transformer}$ 之前，先介绍两个注意力机制在计算机视觉中应用的例子。 $\mathrm{Vision\text{ }Transformer}$ 并不是第一个将注意力机制应用到计算机视觉的领域中去的，其中 $\mathrm{SAGAN}$ 和 $\mathrm{AttnGAN}$ 就早已经在 $\mathrm{GAN}$ 的框架中引入了注意力机制，并且它们大大提高了图像生成的质量。

2.1 Self-Attention GAN

$\mathrm{SAGAN}$ 在 $\mathrm{GAN}$ 的框架中利用自注意力机制来捕获图像特征的长距离依赖关系，使得合成的图像中考量了所有的图像特征信息。 $\mathrm{SAGAN}$ 中自注意力机制的操作原理如下图所示。
给定一个 $3$ 通道的输入特征图 $X=(X^1,X^2,X^3)\in \mathbb{R}^{3\times 3\times 3}$ ，其中 $X^{i}\in \mathbb{R}^{3\times 3}$ ， $i\in\{1,2,3\}$ 。将 $X$ 分别输入到三个不同的 $1\times 1$ 的卷积层中，并生成 $\mathrm{query}$ 特征图 $Q\in \mathbb{R}^{3\times 3\times 3}$ ， $\mathrm{key}$ 特征图 $K\in \mathbb{R}^{3\times 3\times 3}$ 和 $\mathrm{value}$ 特征图 $V\in \mathbb{R}^{3\times 3\times 3}$ 。生成 $Q$ 具体的计算过程为，给定三个卷积核 $W^{q1}$ ， $W^{q2}$ 和 $W^{q3}\in\mathbb{R}^{1\times1\times3}$ ，并用这三个卷积核分别与 $X$ 做卷积运算得到 $Q^1$ ， $Q^2$ 和 $Q^3\in \mathbb{R}^{3 \times 3}$ ，即 $\left\{$

\begin{aligned} Q^{1} & = X * W^{q 1} \\ Q^{2} & = X * W^{q 2} \\ Q^{3} & = X * W^{q 3} \end{aligned}

$\begin{aligned}Q^1&=X * W^{q1}\\Q^2&=X * W^{q2}\\Q^3&=X*W^{q3}\end{aligned}$ \right.

⎩ ⎪ ⎨ ⎪ ⎧ Q^{1} Q^{2} Q^{3} = X * W^{q 1} = X * W^{q 2} = X * W^{q 3}

其中

*

表示卷积运算符号。同理生成

K

和

V

的计算过程与

Q

的计算过程类似。然后再利用

Q

和

K

进行注意力分数的计算得到矩阵

A\in \mathbb{R}^{3 \times 3}

，其中矩阵

A

的元素

a_{ml}

的计算公式为

a_{ml}=Q^m * K^l,\quad m \in \{1,2,3\},l\in \{1,2,3\}

再对矩阵

A

利用

\mathrm{softmax}

函数进行注意力分布的计算得到注意力分布矩阵

S\in \mathbb{R}^{3\times 3}

，其中矩阵

S

的元素

s_{ml}

的计算公式为

s_{ml}=\frac{\exp(a_{ml})}{\sum\limits_{i=j}^{3}\exp(a_{mj})},\quad m \in \{1,2,3\},l\in\{1,2,3\}

最后利用注意力分布矩阵

S

和

\mathrm{value}

特征图

V

得到最后的输出

O=(O^1,O^2,O^3)\in \mathbb{R}^{3\times 3\times 3}

，即

\begin{aligned} O^{1} & = s_{11} \cdot V^{1} + s_{12} \cdot V^{2} + s_{13} \cdot V^{3} \\ O^{2} & = s_{21} \cdot V^{1} + s_{22} \cdot V^{2} + s_{23} \cdot V^{3} \\ O^{3} & = s_{31} \cdot V^{1} + s_{32} \cdot V^{2} + s_{33} \cdot V^{3} \end{aligned}

2.2 AttnGAN

$\mathrm{AttnGAN}$ 通过利用注意力机制来实现多阶段细颗粒度的文本到图像的生成，它可以通过关注自然语言中的一些重要单词来对图像的不同子区域进行合成。比如通过文本“一只鸟有黄色的羽毛和黑色的眼睛”来生成图像时，会对关键词“鸟”，“羽毛”，“眼睛”，“黄色”，“黑色”给予不同的生成权重，并根据这些关键词的引导在图像的不同的子区域中进行细节的丰富。 $\mathrm{AttnGAN}$ 中注意力机制的操作原理如下图所示。
给定输入图像特征向量 $h=(h^1,h^2,h^3,h^4)\in\mathbb{R}^{\hat{D}\times 4}$ 和词特征向量 $e=(e^1,e^2,e^3,e^4)$ ，其中 $h^i\in \mathbb{R}^{\hat{D}\times 1}$ ， $e^i\in \mathbb{R}^{D\times 1}$ ， $i\in \{1,2,3,4\}$ 。首先利用矩阵 $W$ 进行线性变换将词特征空间 $\mathbb{R}^{D}$ 的向量转换成图像特征空间 $\mathbb{R}^{\hat{D}}$ 的向量，则有 $\hat{e}=W\cdot e=(\hat{e}^1,\hat{e}^2,\hat{e}^3,\hat{e}^4)\in \mathbb{R}^{\hat{D}\times 4}$ 然后再利用转换后的词特征 $\hat{e}$ 与图像特征 $h$ 进行注意力分数的计算得到注意力分数矩阵 $S$ ，其中的分量 $s_{ij}$ 的计算公式为 $s_{ij}=(h^i)^{\top}\cdot \hat{e}^j,\quad i\in \{1,2,3,4\},j\in\{1,2,3,4\}$ 再对矩阵 $S$ 利用 $\mathrm{softmax}$ 函数进行注意力分布的计算得到注意力分布矩阵 $\beta\in \mathbb{R}^{4\times 4}$ ，其中矩阵 $\beta$ 的元素 $\beta_{ij}$ 的计算公式为 $\beta_{ij}=\frac{\exp(s_{ij})}{\sum\limits_{k=1}^{3}\exp(s_{ik})},\quad i \in \{1,2,3,4\},l\in\{1,2,3,4\}$ 最后利用注意力分布矩阵 $\beta$ 和图像特征 $h$ 得到最后的输出 $o=(o^1,o^2,o^3,o^4)\in \mathbb{R}^{\hat{D}\times 4}$ ，即 $\left\{$

\begin{aligned} o^{1} & = β_{11} \cdot h^{1} + β_{12} \cdot h^{2} + β_{13} \cdot h^{3} + β_{14} \cdot h^{4} \\ o^{2} & = β_{21} \cdot h^{1} + β_{22} \cdot h^{2} + β_{23} \cdot h^{3} + β_{24} \cdot h^{4} \\ o^{3} & = β_{31} \cdot h^{1} + β_{32} \cdot h^{2} + β_{33} \cdot h^{3} + β_{34} \cdot h^{4} \\ o^{4} & = β_{41} \cdot h^{1} + β_{42} \cdot h^{2} + β_{43} \cdot h^{3} + β_{44} \cdot h^{4} \end{aligned}

$\begin{aligned}o^1&=\beta_{11}\cdot h^1+\beta_{12}\cdot h^2+\beta_{13}\cdot h^3+\beta_{14}\cdot h^4\\o^2&=\beta_{21}\cdot h^1+\beta_{22}\cdot h^2+\beta_{23}\cdot h^3+\beta_{24}\cdot h^4\\o^3&=\beta_{31}\cdot h^1+\beta_{32}\cdot h^2+\beta_{33}\cdot h^3+\beta_{34}\cdot h^4\\o^4&=\beta_{41}\cdot h^1+\beta_{42}\cdot h^2+\beta_{43}\cdot h^3+\beta_{44}\cdot h^4\end{aligned}$ \right.

⎩ ⎪ ⎪ ⎪ ⎪ ⎨ ⎪ ⎪ ⎪ ⎪ ⎧ o^{1} o^{2} o^{3} o^{4} = β_{11} \cdot h^{1} + β_{12} \cdot h^{2} + β_{13} \cdot h^{3} + β_{14} \cdot h^{4} = β_{21} \cdot h^{1} + β_{22} \cdot h^{2} + β_{23} \cdot h^{3} + β_{24} \cdot h^{4} = β_{31} \cdot h^{1} + β_{32} \cdot h^{2} + β_{33} \cdot h^{3} + β_{34} \cdot h^{4} = β_{41} \cdot h^{1} + β_{42} \cdot h^{2} + β_{43} \cdot h^{3} + β_{44} \cdot h^{4}

3 Vision Transformer

本节主要详细介绍 $\mathrm{Vision\text{ }Transformer}$ 的工作原理，3.1节是关于 $\mathrm{Vision\text{ }Transformer}$ 的整体框架，3.2节是关于 $\mathrm{Transformer\text{ }Encoder}$ 的内部操作细节。对于 $\mathrm{Transformer\text{ }Encoder}$ 中 $\mathrm{Multi}$ - $\mathrm{Head\text{ }Attention}$ 的原理本文不会赘述，具体想了解的可以参考上一篇文章《Transformer详解（附代码）》中相关原理的介绍。不难发现，不管是自然语言处理中的 $\mathrm{Transformer}$ ，还是计算机视觉中图像生成的 $\mathrm{SAGAN}$ ，以及文本生成图像的 $\mathrm{AttnGAN}$ ，它们核心模块中注意力机制的主要目的就是求出注意力分布。

3.1 Vision Transformer整体框架

如果下图所示为 $\mathrm{Vision\text{ }Transformer}$ 的整体框架以及相应的训练流程

给定一张图片 $X\in \mathbb{R}^{3n\times 3n}$ ，并将它分割成 $9$ 个 $\mathrm{patch}$ 分别为 $x^1,\cdots,x^9\in\mathbb{R}^{n\times n}$ 。然后再将这个 $9$ 个 $\mathrm{patch}$ 拉平，则有 $x^1,\cdots,x^9\in\mathbb{R}^{n^2}$
利用矩阵 $W\in \mathbb{R}^{l \times n^2}$ 将拉平后的向量 $x^i\in\mathbb{R}^{n^2},i\in\{1,\cdots,9\}$ 经过线性变换得到图像编码向量 $z^i\in \mathbb{R}^{l},i\in\{1,\cdots,9\}$ ，具体的计算公式为 $z^i = W\cdot x^i,\quad i\in\{1,\cdots9\}$
然后将图像编码向量 $z^{i},i\in\{1,\cdot,9\}$ 和类编码向量 $z^0$ 分别与对应的位置编进行加和得到输入编码向量，则有 $z^{i}+p^{i}\in\mathbb{R}^l,\quad i\in\{0,\cdots 9\}$
接着将输入编码向量输入到 $\mathrm{Vision\text{ }Transformer\text{ }Encoder}$ 中得到对应的输出 $o^i\in \mathbb{R}^l,i\in\{0,\cdots,9\}$
最后将类编码向量 $o^0$ 输入全连接神经网络中 $\mathrm{MLP}$ 得到类别预测向量 $\hat{y}\in\mathbb{R}^c$ ，并与真实类别向量 $y\in\mathbb{R}^c$ 计算交叉熵损失得到损失值 $l o s s$ ，利用优化算法更新模型的权重参数

注意事项： 看到这里可能会有一个疑问为什么预测类别的时候只用到了类别编码向量 $o^0$ ， $\mathrm{Vision\text{ }Transformer\text{ }Encoder}$ 其它的输出为什么没有输入到 $\mathrm{MLP}$ 中？为了回答这个问题，我们令函数 $f_0(\cdot)$ 为 $\mathrm{Vision\text{ }Transformer\text{ }Encoder}$ ，则类编码向量 $o^{0}$ 可以表示为 $o^0=f_0(z^0+p^0,\cdots,z^9+p^9)$ 由上公式可以发现，类编码向量 $o^{0}$ 是属于高层特征，其实它综合了所有的图像编码信息，所以可以用它来进行分类，这个可以类比在卷积神经网络中最后的类别输出向量其实就是一层层卷积得到的高层特征。

3.2 Transformer Encoder操作原理

如下图所示分别为 $\mathrm{Vision\text{ }Transformer\text{ }Encoder}$ 模型结构图和原始 $\mathrm{Transformer\text{ }Encoder}$ 的模型结构图。可以直观的发现 $\mathrm{Vision\text{ }Transformer\text{ }Encoder}$ 和 $\mathrm{Transformer\text{ }Encoder}$ 都有层归一化，多头注意力机制，残差连接和线性变换这四个操作，只是在操作顺序有所不同。在以下的 $\mathrm{ \text{ }Transformer}$ 代码实例中，将以下两种 $\mathrm{Encoder}$ 网络结构都进行了实现，可以发现两种网络结构都可以进行很好的训练。
下图左半部分 $\mathrm{Vision\text{ }Transformer\text{ }Encoder}$ 具体的操作流程为

给定输入编码矩阵 $Z\in\mathbb{R}^{l\times n}$ ，首先将其进行层归一化得到 $Z^{\prime}\in\mathbb{R}^{l \times n}$
利用矩阵 $W^{q},W^{k},W^{v}\in \mathbb{R}^{l\times l}$ 对 $Z^{\prime}$ 进行线性变换得到矩阵 $Q,K,W\in\mathbb{R}^{l\times n}$ 具体的计算过程为 $\begin{aligned} Q & = W^{q} \cdot Z^{'} \\ K & = W^{k} \cdot Z^{'} \\ V & = W^{v} \cdot Z^{'} \end{aligned}$ 再将这三个矩阵输入到 $\mathrm{Multi}$ - $\mathrm{Head\text{ }Attention}$ （该原理参考《Transformer详解（附代码）》）中得到矩阵 $Z^{\prime\prime}\in \mathbb{R}^{l \times n}$ 将最原始的输入矩阵 $Z$ 与 $Z^{\prime\prime}$ 进行残差计算得到 $Z+Z^{\prime\prime}\in \mathbb{R}^{l\times n}$
将 $Z+Z^{\prime\prime}$ 进行第二次层归一化得到 $Z^{\prime\prime\prime}\in\mathbb{R}^{l\times n}$ ，然后再将 $Z^{\prime\prime\prime}$ 输入到全连接神经网络中进行线性变换得到 $Z^{\prime\prime\prime\prime}\in\mathbb{R}^{l\times n}$ 。最后将 $Z+Z^{\prime\prime}$ 与 $Z^{\prime\prime\prime\prime}$ 进行残差操作得到该 $\mathrm{Block}$ 的输出 $Z+Z^{\prime\prime}+Z^{\prime\prime\prime\prime}\in\mathbb{R}^{l\times n}$ 。一个 $\mathrm{Encoder}$ 可以将 $N$ 个 $\mathrm{Block}$ 进行堆叠，最后得到的输出为 $O\in\mathbb{R}^{l\times n}$ 。

4 程序代码

$\mathrm{Vision\text{ }Transformer}$ 的代码示例如下所示。该代码是由上一篇《Transformer详解（附代码）》的代码的基础上改编而来。 $\mathrm{Vision\text{ }Transformer}$ 的作者的本意就是想让在 $\mathrm{NLP}$ 中的 $\mathrm{Transformer}$ 模型架构做尽可能少的修改可以直接迁移到 $\mathrm{CV}$ 中，所以以下程序尽可能保持作者的原意，并在代码实现了两种 $\mathrm{Encoder}$ 的网络结构，即3.2节图片所示的两个网络结构，一种是最原始的 $\mathrm{Encoder}$ 网络结构，一种是 $\mathrm{Vision\text{ }Transformer}$ 论文里的 $\mathrm{Encoder}$ 的网络结构。这里需要注意的是， $\mathrm{Vision\text{ }Transformer}$ 里并能没有 $\mathrm{Decoder}$ 模块，所以不需要计算 $\mathrm{Encoder}$ 和 $\mathrm{Decoder}$ 的交叉注意力分布，这就进一步给 $\mathrm{Vision\text{ }Transformer}$ 的编程带来了简便。 $\mathrm{Vision\text{ }Transformer}$ 的开源代码的网址为https://github.com/lucidrains/vit-pytorch/tree/main/vit_pytorch。

import torch
import torch.nn as nn
import os
from einops import rearrange
from einops import repeat
from einops.layers.torch import Rearrange

def inputs_deal(inputs):
	return inputs if isinstance(inputs, tuple) else(inputs, inputs)

class SelfAttention(nn.Module):
	def __init__(self, embed_size, heads):
		super(SelfAttention, self).__init__()
		self.embed_size = embed_size
		self.heads = heads
		self.head_dim = embed_size // heads

		assert (self.head_dim * heads == embed_size), "Embed size needs to be div by heads"

		self.values = nn.Linear(self.head_dim, self.head_dim, bias=False)
		self.keys = nn.Linear(self.head_dim, self.head_dim, bias=False)
		self.queries = nn.Linear(self.head_dim, self.head_dim, bias=False)
		self.fc_out = nn.Linear(heads * self.head_dim, embed_size)

	def forward(self, values, keys, query):
		N =query.shape[0]
		value_len , key_len , query_len = values.shape[1], keys.shape[1], query.shape[1]

		# split embedding into self.heads pieces
		values = values.reshape(N, value_len, self.heads, self.head_dim)
		keys = keys.reshape(N, key_len, self.heads, self.head_dim)
		queries = query.reshape(N, query_len, self.heads, self.head_dim)
		
		values = self.values(values)
		keys = self.keys(keys)
		queries = self.queries(queries)

		energy = torch.einsum("nqhd,nkhd->nhqk", queries, keys)
		# queries shape: (N, query_len, heads, heads_dim)
		# keys shape : (N, key_len, heads, heads_dim)
		# energy shape: (N, heads, query_len, key_len)

		attention = torch.softmax(energy/ (self.embed_size ** (1/2)), dim=3)

		out = torch.einsum("nhql, nlhd->nqhd", [attention, values]).reshape(N, query_len, self.heads*self.head_dim)
		# attention shape: (N, heads, query_len, key_len)
		# values shape: (N, value_len, heads, heads_dim)
		# (N, query_len, heads, head_dim)

		out = self.fc_out(out)
		return out


class TransformerBlock(nn.Module):
	def __init__(self, embed_size, heads, dropout, forward_expansion):
		super(TransformerBlock, self).__init__()
		self.attention = SelfAttention(embed_size, heads)
		self.norm = nn.LayerNorm(embed_size)

		self.feed_forward = nn.Sequential(
			nn.Linear(embed_size, forward_expansion*embed_size),
			nn.ReLU(),
			nn.Linear(forward_expansion*embed_size, embed_size)
		)
		self.dropout = nn.Dropout(dropout)


	def forward(self, value, key, query, x, type_mode):
		if type_mode == 'original':
			attention = self.attention(value, key, query)
			x = self.dropout(self.norm(attention + x))
			forward = self.feed_forward(x)
			out = self.dropout(self.norm(forward + x))
			return out
		else:
			attention = self.attention(self.norm(value), self.norm(key), self.norm(query))
			x =self.dropout(attention + x)
			forward = self.feed_forward(self.norm(x))
			out = self.dropout(forward + x)
			return out

class TransformerEncoder(nn.Module):
	def __init__(
			self,
			embed_size,
			num_layers,
			heads,
			forward_expansion,
			dropout = 0,
			type_mode = 'original'
		):
		super(TransformerEncoder, self).__init__()
		self.embed_size = embed_size
		self.type_mode = type_mode
		self.Query_Key_Value = nn.Linear(embed_size, embed_size * 3, bias = False)

		self.layers = nn.ModuleList(
			[
				TransformerBlock(
					embed_size,
					heads,
					dropout=dropout,
					forward_expansion=forward_expansion,
					)
				for _ in range(num_layers)]
		)
		self.dropout = nn.Dropout(dropout)

	def forward(self, x):
		for layer in self.layers:
			QKV_list = self.Query_Key_Value(x).chunk(3, dim = -1)
			x = layer(QKV_list[0], QKV_list[1], QKV_list[2], x, self.type_mode)
		return x

class VisionTransformer(nn.Module):
	def __init__(self, 
				image_size, 
				patch_size, 
				num_classes, 
				embed_size, 
				num_layers, 
				heads, 
				mlp_dim, 
				pool = 'cls',
				channels = 3,
				dropout = 0,
				emb_dropout = 0.1,
				type_mode = 'vit'):
		super(VisionTransformer, self).__init__()
		img_h, img_w = inputs_deal(image_size)
		patch_h, patch_w = inputs_deal(patch_size)
		
		assert img_h % patch_h == 0 and img_w % patch_w == 0, 'Img dimensions can be divisible by the patch dimensions'

		num_patches = (img_h // patch_h) * (img_w // patch_w)

		patch_size = channels * patch_h * patch_w

		self.patch_embedding = nn.Sequential(
			Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1 = patch_h, p2=patch_w),
			nn.Linear(patch_size, embed_size, bias=False)
		)


		self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, embed_size))
		self.cls_token = nn.Parameter(torch.randn(1, 1, embed_size))
		self.dropout = nn.Dropout(emb_dropout)



		self.transformer = TransformerEncoder(embed_size, 
									num_layers, 
									heads, 
									mlp_dim,
									dropout)
		self.pool = pool
		self.to_latent = nn.Identity()

		self.mlp_head = nn.Sequential(
			nn.LayerNorm(embed_size),
			nn.Linear(embed_size, num_classes)
		)

	def forward(self, img):
		x = self.patch_embedding(img)
		b, n, _ = x.shape
		cls_tokens = repeat(self.cls_token, '() n d ->b n d', b = b)
		x = torch.cat((cls_tokens, x), dim = 1)
		x += self.pos_embedding[:, :(n + 1)]
		x = self.dropout(x)
		x = self.transformer(x)
		x = x.mean(dim = 1) if self.pool == 'mean' else x[:, 0]
		x = self.to_latent(x)
		return self.mlp_head(x)


if __name__ == '__main__':
	vit = VisionTransformer(
			image_size = 256,
			patch_size = 16,
			num_classes = 10,
			embed_size = 256,
			num_layers = 6,
			heads = 8,
			mlp_dim = 512,
			dropout = 0.1,
			emb_dropout = 0.1
		)
	img = torch.randn(3, 3, 256, 256)
	pred = vit(img)
	print(pred)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191

以下代码是利用 $\mathrm{Vision \text{ }Transformer}$ 网络结构训练一个分类 $\mathrm{mnist}$ 数据集的主程序代码。

from torchvision import datasets, transforms
from torch.utils.data import DataLoader, Dataset
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import VIT
import os

        
def train():
    batch_size = 4
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    epoches = 20
    mnist_train = datasets.MNIST("mnist-data", train=True, download=True, transform=transforms.ToTensor())
    train_loader = torch.utils.data.DataLoader(mnist_train, batch_size= batch_size, shuffle=True)
    mnist_model = VIT.VisionTransformer(
    	image_size = 28,
    	patch_size = 7,
    	num_classes = 10,
    	channels = 1,
    	embed_size = 512,
    	num_layers = 1,
    	heads = 2,
    	mlp_dim =1024,
    	dropout = 0,
    	emb_dropout = 0)
    loss_fn = nn.CrossEntropyLoss()
    mnist_model = mnist_model.to(device)
    opitimizer = optim.Adam(mnist_model.parameters(), lr=0.00001)
    mnist_model.train()
    for epoch in range(epoches):
    	total_loss = 0 
    	corrects = 0 
    	num = 0
    	for batch_X, batch_Y in train_loader:
    		batch_X, batch_Y = batch_X.to(device), batch_Y.to(device)
    		opitimizer.zero_grad()
    		outputs = mnist_model(batch_X)
    		_, pred = torch.max(outputs.data, 1)
    		loss = loss_fn(outputs, batch_Y)
    		loss.backward()
    		opitimizer.step()
    		total_loss += loss.item()
    		corrects = torch.sum(pred == batch_Y.data)
    		num += batch_size
    		print(epoch, total_loss/float(num), corrects.item()/float(batch_size))

if __name__ == '__main__':
	train()

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51

训练的过程如下所示，可以发现损失函数可以稳定下降。但是训练一个 $\mathrm{Vision \text{ }Transformer}$ 模型真的是很烧硬件，跟训练一个普通的 $\mathrm{CNN}$ 模型相比，训练一个 $\mathrm{Vision \text{ }Transformer}$ 模型更加耗时耗力。

声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/w/酷酷是懒虫/article/detail/764329