如何在 PyTorch 中初始化权重?
How to initialize weights in PyTorch?
如何初始化网络的权重和偏差(例如通过 He 或 Xavier 初始化)?
单层
要初始化单个层的权重,请使用 torch.nn.init
中的函数。例如:
conv1 = torch.nn.Conv2d(...)
torch.nn.init.xavier_uniform(conv1.weight)
或者,您可以通过写入 conv1.weight.data
(即 torch.Tensor
)来修改参数。示例:
conv1.weight.data.fill_(0.01)
同样适用于偏见:
conv1.bias.data.fill_(0.01)
nn.Sequential
或自定义 nn.Module
将初始化函数传递给 torch.nn.Module.apply
。它将递归地初始化整个nn.Module
中的权重。
apply(fn): Applies fn
recursively to every submodule (as returned by .children()
) as well as self. Typical use includes initializing the parameters of a model (see also torch-nn-init).
示例:
def init_weights(m):
if isinstance(m, nn.Linear):
torch.nn.init.xavier_uniform(m.weight)
m.bias.data.fill_(0.01)
net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))
net.apply(init_weights)
抱歉来晚了,希望我的回答对您有所帮助。
要使用 normal distribution
初始化权重,请使用:
torch.nn.init.normal_(tensor, mean=0, std=1)
或者使用 constant distribution
写:
torch.nn.init.constant_(tensor, value)
或者使用 uniform distribution
:
torch.nn.init.uniform_(tensor, a=0, b=1) # a: lower_bound, b: upper_bound
您可以查看其他初始化张量的方法here
import torch.nn as nn
# a simple network
rand_net = nn.Sequential(nn.Linear(in_features, h_size),
nn.BatchNorm1d(h_size),
nn.ReLU(),
nn.Linear(h_size, h_size),
nn.BatchNorm1d(h_size),
nn.ReLU(),
nn.Linear(h_size, 1),
nn.ReLU())
# initialization function, first checks the module type,
# then applies the desired changes to the weights
def init_normal(m):
if type(m) == nn.Linear:
nn.init.uniform_(m.weight)
# use the modules apply function to recursively apply the initialization
rand_net.apply(init_normal)
我们使用相同的神经网络 (NN) 架构比较不同的权重初始化模式。
全零或全一
如果遵循Occam's razor的原则,您可能会认为将所有权重设置为 0 或 1 将是最佳解决方案。事实并非如此。
在每个权重相同的情况下,每一层的所有神经元都产生相同的输出。这使得很难决定调整哪些权重。
# initialize two NN's with 0 and 1 constant weights
model_0 = Net(constant_weight=0)
model_1 = Net(constant_weight=1)
- 2个epoch后:
Validation Accuracy
9.625% -- All Zeros
10.050% -- All Ones
Training Loss
2.304 -- All Zeros
1552.281 -- All Ones
统一初始化
A uniform distribution 从一组数字中选出任意数字的概率相等。
让我们看看使用统一权重初始化的神经网络训练效果如何,其中 low=0.0
和 high=1.0
。
下面,我们将看到另一种方式(除了在 Net class 代码中)来初始化网络的权重。要在模型定义之外定义权重,我们可以:
- Define a function that assigns weights by the type of network layer, then
- Apply those weights to an initialized model using
model.apply(fn)
, which applies a function to each model layer.
# takes in a module and applies the specified weight initialization
def weights_init_uniform(m):
classname = m.__class__.__name__
# for every Linear layer in a model..
if classname.find('Linear') != -1:
# apply a uniform distribution to the weights and a bias=0
m.weight.data.uniform_(0.0, 1.0)
m.bias.data.fill_(0)
model_uniform = Net()
model_uniform.apply(weights_init_uniform)
- 2个epoch后:
Validation Accuracy
36.667% -- Uniform Weights
Training Loss
3.208 -- Uniform Weights
设置权重的一般规则
在神经网络中设置权重的一般规则是将它们设置为接近于零但又不会太小。
Good practice is to start your weights in the range of [-y, y] where y=1/sqrt(n)
(n is the number of inputs to a given neuron).
# takes in a module and applies the specified weight initialization
def weights_init_uniform_rule(m):
classname = m.__class__.__name__
# for every Linear layer in a model..
if classname.find('Linear') != -1:
# get the number of the inputs
n = m.in_features
y = 1.0/np.sqrt(n)
m.weight.data.uniform_(-y, y)
m.bias.data.fill_(0)
# create a new model with these weights
model_rule = Net()
model_rule.apply(weights_init_uniform_rule)
下面我们比较 NN 的性能,使用均匀分布 [-0.5,0.5) 初始化的权重与使用 一般规则
初始化权重的神经网络的性能
- 2个epoch后:
Validation Accuracy
75.817% -- Centered Weights [-0.5, 0.5)
85.208% -- General Rule [-y, y)
Training Loss
0.705 -- Centered Weights [-0.5, 0.5)
0.469 -- General Rule [-y, y)
初始化权重的正态分布
The normal distribution should have a mean of 0 and a standard deviation of y=1/sqrt(n)
, where n is the number of inputs to NN
## takes in a module and applies the specified weight initialization
def weights_init_normal(m):
'''Takes in a module and initializes all linear layers with weight
values taken from a normal distribution.'''
classname = m.__class__.__name__
# for every Linear layer in a model
if classname.find('Linear') != -1:
y = m.in_features
# m.weight.data shoud be taken from a normal distribution
m.weight.data.normal_(0.0,1/np.sqrt(y))
# m.bias.data should be 0
m.bias.data.fill_(0)
下面我们展示了两个 NN 的性能,一个使用 uniform-distribution 初始化,另一个使用 normal-distribution
- 2个epoch后:
Validation Accuracy
85.775% -- Uniform Rule [-y, y)
84.717% -- Normal Distribution
Training Loss
0.329 -- Uniform Rule [-y, y)
0.443 -- Normal Distribution
如果您看到弃用警告 (@Fábio Perez)...
def init_weights(m):
if type(m) == nn.Linear:
torch.nn.init.xavier_uniform_(m.weight)
m.bias.data.fill_(0.01)
net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))
net.apply(init_weights)
要初始化图层,您通常不需要执行任何操作。
PyTorch 会为你做这件事。如果您考虑一下,这很有意义。当 PyTorch 可以按照最新趋势进行初始化时,我们为什么要初始化层。
例如检查 Linear layer。
在 __init__
方法中它将调用 Kaiming He 初始化函数。
def reset_parameters(self):
init.kaiming_uniform_(self.weight, a=math.sqrt(3))
if self.bias is not None:
fan_in, _ = init._calculate_fan_in_and_fan_out(self.weight)
bound = 1 / math.sqrt(fan_in)
init.uniform_(self.bias, -bound, bound)
其他图层类型也是如此。对于 conv2d
例如检查 here.
注意:正确初始化的好处是训练速度更快。
如果您的问题需要特殊初始化,您可以在之后进行。
遍历参数
如果您不能使用 apply
,例如如果模型没有直接实现 Sequential
:
所有人都一样
# see UNet at https://github.com/milesial/Pytorch-UNet/tree/master/unet
def init_all(model, init_func, *params, **kwargs):
for p in model.parameters():
init_func(p, *params, **kwargs)
model = UNet(3, 10)
init_all(model, torch.nn.init.normal_, mean=0., std=1)
# or
init_all(model, torch.nn.init.constant_, 1.)
取决于形状
def init_all(model, init_funcs):
for p in model.parameters():
init_func = init_funcs.get(len(p.shape), init_funcs["default"])
init_func(p)
model = UNet(3, 10)
init_funcs = {
1: lambda x: torch.nn.init.normal_(x, mean=0., std=1.), # can be bias
2: lambda x: torch.nn.init.xavier_normal_(x, gain=1.), # can be weight
3: lambda x: torch.nn.init.xavier_uniform_(x, gain=1.), # can be conv1D filter
4: lambda x: torch.nn.init.xavier_uniform_(x, gain=1.), # can be conv2D filter
"default": lambda x: torch.nn.init.constant(x, 1.), # everything else
}
init_all(model, init_funcs)
您可以尝试使用 torch.nn.init.constant_(x, len(x.shape))
检查它们是否已正确初始化:
init_funcs = {
"default": lambda x: torch.nn.init.constant_(x, len(x.shape))
}
如果您想要更多的灵活性,您也可以手动设置权重。
假设你输入了所有的:
import torch
import torch.nn as nn
input = torch.ones((8, 8))
print(input)
tensor([[1., 1., 1., 1., 1., 1., 1., 1.],
[1., 1., 1., 1., 1., 1., 1., 1.],
[1., 1., 1., 1., 1., 1., 1., 1.],
[1., 1., 1., 1., 1., 1., 1., 1.],
[1., 1., 1., 1., 1., 1., 1., 1.],
[1., 1., 1., 1., 1., 1., 1., 1.],
[1., 1., 1., 1., 1., 1., 1., 1.],
[1., 1., 1., 1., 1., 1., 1., 1.]])
并且您想制作一个没有偏差的致密层(以便我们可以想象):
d = nn.Linear(8, 8, bias=False)
将所有权重设置为 0.5(或其他任何值):
d.weight.data = torch.full((8, 8), 0.5)
print(d.weight.data)
权重:
Out[14]:
tensor([[0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000],
[0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000],
[0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000],
[0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000],
[0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000],
[0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000],
[0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000],
[0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000]])
你所有的权重现在都是 0.5。通过以下方式传递数据:
d(input)
Out[13]:
tensor([[4., 4., 4., 4., 4., 4., 4., 4.],
[4., 4., 4., 4., 4., 4., 4., 4.],
[4., 4., 4., 4., 4., 4., 4., 4.],
[4., 4., 4., 4., 4., 4., 4., 4.],
[4., 4., 4., 4., 4., 4., 4., 4.],
[4., 4., 4., 4., 4., 4., 4., 4.],
[4., 4., 4., 4., 4., 4., 4., 4.],
[4., 4., 4., 4., 4., 4., 4., 4.]], grad_fn=<MmBackward>)
请记住,每个神经元接收 8 个输入,所有输入的权重均为 0.5,值为 1(且无偏差),因此每个神经元的总和为 4。
因为到目前为止我还没有足够的声誉,我无法在
下添加评论
the answer posted by prosti in Jun 26 '19 at 13:16.
def reset_parameters(self):
init.kaiming_uniform_(self.weight, a=math.sqrt(3))
if self.bias is not None:
fan_in, _ = init._calculate_fan_in_and_fan_out(self.weight)
bound = 1 / math.sqrt(fan_in)
init.uniform_(self.bias, -bound, bound)
但我想指出,实际上我们知道 Kaiming He、Delving Deep into Rectifiers: Surpassing Human-Level Performance on ImageNet 论文中的一些假设分类,不合适,虽然看起来刻意设计的初始化方法在实践中很成功。
例如,在 Backward Propagation Case 的小节中,他们假设 $w_l$ 和 $\delta y_l$ 彼此独立.但众所周知,以得分图$\delta y^L_i$为例,往往是$y_i-softmax(y^L_i)=y_i-softmax( w^L_ix^L_i)$如果我们使用典型的交叉熵损失函数objective.
所以我认为 He's Initialization 运作良好的真正根本原因仍有待阐明。因为每个人都见证了它在促进深度学习训练方面的力量。
这是更好的方法,只需传递整个模型
import torch.nn as nn
def initialize_weights(model):
# Initializes weights according to the DCGAN paper
for m in model.modules():
if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d, nn.BatchNorm2d)):
nn.init.normal_(m.weight.data, 0.0, 0.02)
# if you also want for linear layers ,add one more elif condition
如何初始化网络的权重和偏差(例如通过 He 或 Xavier 初始化)?
单层
要初始化单个层的权重,请使用 torch.nn.init
中的函数。例如:
conv1 = torch.nn.Conv2d(...)
torch.nn.init.xavier_uniform(conv1.weight)
或者,您可以通过写入 conv1.weight.data
(即 torch.Tensor
)来修改参数。示例:
conv1.weight.data.fill_(0.01)
同样适用于偏见:
conv1.bias.data.fill_(0.01)
nn.Sequential
或自定义 nn.Module
将初始化函数传递给 torch.nn.Module.apply
。它将递归地初始化整个nn.Module
中的权重。
apply(fn): Applies
fn
recursively to every submodule (as returned by.children()
) as well as self. Typical use includes initializing the parameters of a model (see also torch-nn-init).
示例:
def init_weights(m):
if isinstance(m, nn.Linear):
torch.nn.init.xavier_uniform(m.weight)
m.bias.data.fill_(0.01)
net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))
net.apply(init_weights)
抱歉来晚了,希望我的回答对您有所帮助。
要使用 normal distribution
初始化权重,请使用:
torch.nn.init.normal_(tensor, mean=0, std=1)
或者使用 constant distribution
写:
torch.nn.init.constant_(tensor, value)
或者使用 uniform distribution
:
torch.nn.init.uniform_(tensor, a=0, b=1) # a: lower_bound, b: upper_bound
您可以查看其他初始化张量的方法here
import torch.nn as nn
# a simple network
rand_net = nn.Sequential(nn.Linear(in_features, h_size),
nn.BatchNorm1d(h_size),
nn.ReLU(),
nn.Linear(h_size, h_size),
nn.BatchNorm1d(h_size),
nn.ReLU(),
nn.Linear(h_size, 1),
nn.ReLU())
# initialization function, first checks the module type,
# then applies the desired changes to the weights
def init_normal(m):
if type(m) == nn.Linear:
nn.init.uniform_(m.weight)
# use the modules apply function to recursively apply the initialization
rand_net.apply(init_normal)
我们使用相同的神经网络 (NN) 架构比较不同的权重初始化模式。
全零或全一
如果遵循Occam's razor的原则,您可能会认为将所有权重设置为 0 或 1 将是最佳解决方案。事实并非如此。
在每个权重相同的情况下,每一层的所有神经元都产生相同的输出。这使得很难决定调整哪些权重。
# initialize two NN's with 0 and 1 constant weights
model_0 = Net(constant_weight=0)
model_1 = Net(constant_weight=1)
- 2个epoch后:
Validation Accuracy
9.625% -- All Zeros
10.050% -- All Ones
Training Loss
2.304 -- All Zeros
1552.281 -- All Ones
统一初始化
A uniform distribution 从一组数字中选出任意数字的概率相等。
让我们看看使用统一权重初始化的神经网络训练效果如何,其中 low=0.0
和 high=1.0
。
下面,我们将看到另一种方式(除了在 Net class 代码中)来初始化网络的权重。要在模型定义之外定义权重,我们可以:
- Define a function that assigns weights by the type of network layer, then
- Apply those weights to an initialized model using
model.apply(fn)
, which applies a function to each model layer.
# takes in a module and applies the specified weight initialization
def weights_init_uniform(m):
classname = m.__class__.__name__
# for every Linear layer in a model..
if classname.find('Linear') != -1:
# apply a uniform distribution to the weights and a bias=0
m.weight.data.uniform_(0.0, 1.0)
m.bias.data.fill_(0)
model_uniform = Net()
model_uniform.apply(weights_init_uniform)
- 2个epoch后:
Validation Accuracy
36.667% -- Uniform Weights
Training Loss
3.208 -- Uniform Weights
设置权重的一般规则
在神经网络中设置权重的一般规则是将它们设置为接近于零但又不会太小。
Good practice is to start your weights in the range of [-y, y] where
y=1/sqrt(n)
(n is the number of inputs to a given neuron).
# takes in a module and applies the specified weight initialization
def weights_init_uniform_rule(m):
classname = m.__class__.__name__
# for every Linear layer in a model..
if classname.find('Linear') != -1:
# get the number of the inputs
n = m.in_features
y = 1.0/np.sqrt(n)
m.weight.data.uniform_(-y, y)
m.bias.data.fill_(0)
# create a new model with these weights
model_rule = Net()
model_rule.apply(weights_init_uniform_rule)
下面我们比较 NN 的性能,使用均匀分布 [-0.5,0.5) 初始化的权重与使用 一般规则
初始化权重的神经网络的性能- 2个epoch后:
Validation Accuracy
75.817% -- Centered Weights [-0.5, 0.5)
85.208% -- General Rule [-y, y)
Training Loss
0.705 -- Centered Weights [-0.5, 0.5)
0.469 -- General Rule [-y, y)
初始化权重的正态分布
The normal distribution should have a mean of 0 and a standard deviation of
y=1/sqrt(n)
, where n is the number of inputs to NN
## takes in a module and applies the specified weight initialization
def weights_init_normal(m):
'''Takes in a module and initializes all linear layers with weight
values taken from a normal distribution.'''
classname = m.__class__.__name__
# for every Linear layer in a model
if classname.find('Linear') != -1:
y = m.in_features
# m.weight.data shoud be taken from a normal distribution
m.weight.data.normal_(0.0,1/np.sqrt(y))
# m.bias.data should be 0
m.bias.data.fill_(0)
下面我们展示了两个 NN 的性能,一个使用 uniform-distribution 初始化,另一个使用 normal-distribution
- 2个epoch后:
Validation Accuracy
85.775% -- Uniform Rule [-y, y)
84.717% -- Normal Distribution
Training Loss
0.329 -- Uniform Rule [-y, y)
0.443 -- Normal Distribution
如果您看到弃用警告 (@Fábio Perez)...
def init_weights(m):
if type(m) == nn.Linear:
torch.nn.init.xavier_uniform_(m.weight)
m.bias.data.fill_(0.01)
net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))
net.apply(init_weights)
要初始化图层,您通常不需要执行任何操作。 PyTorch 会为你做这件事。如果您考虑一下,这很有意义。当 PyTorch 可以按照最新趋势进行初始化时,我们为什么要初始化层。
例如检查 Linear layer。
在 __init__
方法中它将调用 Kaiming He 初始化函数。
def reset_parameters(self):
init.kaiming_uniform_(self.weight, a=math.sqrt(3))
if self.bias is not None:
fan_in, _ = init._calculate_fan_in_and_fan_out(self.weight)
bound = 1 / math.sqrt(fan_in)
init.uniform_(self.bias, -bound, bound)
其他图层类型也是如此。对于 conv2d
例如检查 here.
注意:正确初始化的好处是训练速度更快。 如果您的问题需要特殊初始化,您可以在之后进行。
遍历参数
如果您不能使用 apply
,例如如果模型没有直接实现 Sequential
:
所有人都一样
# see UNet at https://github.com/milesial/Pytorch-UNet/tree/master/unet
def init_all(model, init_func, *params, **kwargs):
for p in model.parameters():
init_func(p, *params, **kwargs)
model = UNet(3, 10)
init_all(model, torch.nn.init.normal_, mean=0., std=1)
# or
init_all(model, torch.nn.init.constant_, 1.)
取决于形状
def init_all(model, init_funcs):
for p in model.parameters():
init_func = init_funcs.get(len(p.shape), init_funcs["default"])
init_func(p)
model = UNet(3, 10)
init_funcs = {
1: lambda x: torch.nn.init.normal_(x, mean=0., std=1.), # can be bias
2: lambda x: torch.nn.init.xavier_normal_(x, gain=1.), # can be weight
3: lambda x: torch.nn.init.xavier_uniform_(x, gain=1.), # can be conv1D filter
4: lambda x: torch.nn.init.xavier_uniform_(x, gain=1.), # can be conv2D filter
"default": lambda x: torch.nn.init.constant(x, 1.), # everything else
}
init_all(model, init_funcs)
您可以尝试使用 torch.nn.init.constant_(x, len(x.shape))
检查它们是否已正确初始化:
init_funcs = {
"default": lambda x: torch.nn.init.constant_(x, len(x.shape))
}
如果您想要更多的灵活性,您也可以手动设置权重。
假设你输入了所有的:
import torch
import torch.nn as nn
input = torch.ones((8, 8))
print(input)
tensor([[1., 1., 1., 1., 1., 1., 1., 1.],
[1., 1., 1., 1., 1., 1., 1., 1.],
[1., 1., 1., 1., 1., 1., 1., 1.],
[1., 1., 1., 1., 1., 1., 1., 1.],
[1., 1., 1., 1., 1., 1., 1., 1.],
[1., 1., 1., 1., 1., 1., 1., 1.],
[1., 1., 1., 1., 1., 1., 1., 1.],
[1., 1., 1., 1., 1., 1., 1., 1.]])
并且您想制作一个没有偏差的致密层(以便我们可以想象):
d = nn.Linear(8, 8, bias=False)
将所有权重设置为 0.5(或其他任何值):
d.weight.data = torch.full((8, 8), 0.5)
print(d.weight.data)
权重:
Out[14]:
tensor([[0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000],
[0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000],
[0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000],
[0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000],
[0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000],
[0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000],
[0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000],
[0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000, 0.5000]])
你所有的权重现在都是 0.5。通过以下方式传递数据:
d(input)
Out[13]:
tensor([[4., 4., 4., 4., 4., 4., 4., 4.],
[4., 4., 4., 4., 4., 4., 4., 4.],
[4., 4., 4., 4., 4., 4., 4., 4.],
[4., 4., 4., 4., 4., 4., 4., 4.],
[4., 4., 4., 4., 4., 4., 4., 4.],
[4., 4., 4., 4., 4., 4., 4., 4.],
[4., 4., 4., 4., 4., 4., 4., 4.],
[4., 4., 4., 4., 4., 4., 4., 4.]], grad_fn=<MmBackward>)
请记住,每个神经元接收 8 个输入,所有输入的权重均为 0.5,值为 1(且无偏差),因此每个神经元的总和为 4。
因为到目前为止我还没有足够的声誉,我无法在
下添加评论the answer posted by prosti in Jun 26 '19 at 13:16.
def reset_parameters(self):
init.kaiming_uniform_(self.weight, a=math.sqrt(3))
if self.bias is not None:
fan_in, _ = init._calculate_fan_in_and_fan_out(self.weight)
bound = 1 / math.sqrt(fan_in)
init.uniform_(self.bias, -bound, bound)
但我想指出,实际上我们知道 Kaiming He、Delving Deep into Rectifiers: Surpassing Human-Level Performance on ImageNet 论文中的一些假设分类,不合适,虽然看起来刻意设计的初始化方法在实践中很成功。
例如,在 Backward Propagation Case 的小节中,他们假设 $w_l$ 和 $\delta y_l$ 彼此独立.但众所周知,以得分图$\delta y^L_i$为例,往往是$y_i-softmax(y^L_i)=y_i-softmax( w^L_ix^L_i)$如果我们使用典型的交叉熵损失函数objective.
所以我认为 He's Initialization 运作良好的真正根本原因仍有待阐明。因为每个人都见证了它在促进深度学习训练方面的力量。
这是更好的方法,只需传递整个模型
import torch.nn as nn
def initialize_weights(model):
# Initializes weights according to the DCGAN paper
for m in model.modules():
if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d, nn.BatchNorm2d)):
nn.init.normal_(m.weight.data, 0.0, 0.02)
# if you also want for linear layers ,add one more elif condition