torch.nn.Linear的维度变换过程详解(有图有公式有代码)

学技术的大胜嗷 2024-09-10 13:31:02 阅读 72

当初在学习nn.Linear时了解到的博客都是关于一维变换的,比如输入3通道,输出6通道;又比如得到(3,4,4)的特征图,需要进行拉平为(48,)的向量,然后通过nn.Linear(48,10)得到10个输出(分类任务很常见)。

<code>nn.Linear除了可以进行分类,主要的作用就是改变维度便于下一个卷积层或线形层的输入。

但是在实际代码中,nn.Linear的输入往往都是多维数据,一样可以正常输出。所以经过查阅手册和各个帖子,给出了自己的理解,作为笔记。

目录

一、nn.Linear函数用法

二、维度变换过程

三、全连接层的参数量与计算量

一、nn.Linear函数用法

nn.Linear 是 PyTorch 框架中的一个模块,用于实现线性层,也就是全连接层。线性层是神经网络中的基本构件,它执行一个基于矩阵乘法的线性变换,通常用于将输入数据转换为输出数据。

参数介绍:

in_features:输入特征的数量。out_features:输出特征的数量。bias:一个布尔值,指示是否使用偏置项(默认为True)。

import torch

import torch.nn as nn

# 定义输入特征的尺寸

input_height, input_width = 4, 4

# 定义输入通道数

input_channels = 3

# 定义输出节点数

output_nodes = 5

# 创建一个随机的输入特征图,维度为[2,3,4,4]

input_data = torch.randn(2, input_channels, input_height, input_width)

# 创建一个全连接层,4 -> 5

linear_layer = nn.Linear(input_data.size(-1), output_nodes)

# 应用全连接层

output = linear_layer(input_data)

# 输出的尺寸将是 [2,3,4,5]

print("Output shape:", output.shape)

可以看到[2,3,4,4]维度的数据经过<code>nn.Linear得到了[2,3,4,5]的数据,确实可以计算多维度。

二、维度变换过程

我查了pytorch的手册,如下:

首先<code>通过公式可以看到nn.Linear是通过一个权重矩阵来实现维度的变化的。x是输入,A是权重矩阵,x与经过转置的权重矩阵A进行矩阵乘法,最后加上偏置项。

其次nn.Linear的输入是不限制维度的,可以看到括号中的*,其中 * 表示任意数量的附加维度,包括为空(即常见的数据拉平后只剩一个维度)。

权重矩阵维度为(out,in),但是nn.Linear函数的用法是nn.Linear(in,out)。

最终输出的结果是(*,out)。

我画了个计算维度变换图,如下:

假设输入的数据维度为[32,3,4],通过<code>nn.Linear(4,2)得到[32,3,2]。这里取消偏置项

由于在手册中权重矩阵的维度是(out,in),那么而经过转置之后就是(in,out)也就是图中的(4,2)。|

最终得到(1,2)形状的输出,准确的来说,是将(4,)形状变为(2,)。

图中可以看到单个权重矩阵有8个参数,好像不多,为什么其他帖子中都说全连接层的参数量很大呢?

三、全连接层的参数量与计算量

这一章用代码输出数据来论证,还是以输入的数据维度为[32,3,4],通过nn.Linear(4,2)得到[32,3,2]为例,在给出代码之前先猜一下两个问题。

1.  这个过程的参数量是多少?

2. 这个过程的计算量是多少?

我在很多帖子上看到说全连接层参数量很大等等结论,于是我一开始以为参数量是32*3*4*2=768,计算量也是这么多。但是实际情况并不是(他们说的是维度拉平后再输入的情况),代码如下:

import torch

import torch.nn as nn

from thop import profile

from thop import clever_format

class MyModel(nn.Module):

def __init__(self, input_k, output_nodes):

super(MyModel, self).__init__()

# 全连接层

self.linear = nn.Linear(input_k, output_nodes, bias=False)

def forward(self, x):

# 应用全连接层

x = self.linear(x)

return x

# 定义输入特征的尺寸

input_k = 4

# 定义输入通道数

input_channels = 3

# 定义输出节点数

output_nodes = 2

# 创建一个随机的输入特征图,维度为[32,3,4]

input_data = torch.randn(32, input_channels, input_k)

# 创建一个全连接层,4 -> 2

model = MyModel(input_data.size(-1), output_nodes)

# 应用全连接层

output = model(input_data)

# 输出的尺寸将是 [32,3,2]

print("Output shape:", output.shape)

# 定义一个函数来计算模型的参数量

def count_parameters(model):

return sum(p.numel() for p in model.parameters())

# 计算并打印模型的参数量

total_params = count_parameters(model)

print("Total parameters:", total_params)

# 使用 thop 计算 FLOPs

flops, params = profile(model.to('cuda'), (input_data.to('cuda'), ), verbose=False)

# flops 已经是浮点数

print('Total GFLOPS: %s' % flops, 'Total params: %s' % params)

关于代码我解释一下,定义了一个只包含一个线性层的模型,便于计算参数量计算量

可以看到输出的形状只改变了最后一个维度,从[32,3,4]变为[32,3,2]。

但是参数量却等于8,只等于一个权重矩阵,难道是最后一个维度共享权重矩阵么?

于是我利用thop库得到计算量(也就是前向过程计算了多少次),发现是768,正好等于32*3*4*2

现在来分析一下,目前来看<code>nn.Linear只会改变数据最后一个维度的大小。那么就不会对每个样本的所有维度都分配单独的权重,这就是手册中权重矩阵的维度是(out,in)的原因,原来一切早已注定,只是官方没解释太详细。

所以现在是数据的最后维度都是共享权重,权重参数量为4*2=8,所以参数量总和是8。

每更新一次权重参数就算8次计算(每个权重矩阵有8个参数),也就是说遍历完输入数据的维度需要32*3次,那么32*3*(4*2)=768

计算量为768也验证了共享权重的猜想。

现在回头来看手册中的内容,就理解其中的内容了。

不足之处请大佬指出!



声明

本文内容仅代表作者观点,或转载于其他网站,本站不以此文作为商业用途
如有涉及侵权,请联系本站进行删除
转载本站原创文章,请注明来源及作者。