在某些情况下,我们希望将现有模型的激活函数统一替换成某种新型激活函数试验一下效果,但是直接去修改原来的模型文件很麻烦,更好的办法是对生成的模型进行统一替换。
def convert_relu_to_softplus(model):
for child_name, child in model.named_children():
if isinstance(child, nn.ReLU):
setattr(model, child_name, nn.Softplus())
else:
convert_relu_to_softplus(child)
对于这样定义的模型将不起作用:因为没有在初始化函数中注册。
class MyModel(nn.Module):
def forward(self, x):
...
x = F.relu(x)
...
但是,它适用于以下模型:
class MyModel(nn.Module):
def __init__(self):
super().__init__()
...
self.relu = nn.ReLU()
...
def forward(self, x):
...
x = self.relu(x)
...
参考:
https://discuss.pytorch.org/t/how-to-replace-all-relu-activations-in-a-pretrained-network/31591/8