在预训练模型的开头添加归一化层

时间:2020-12-29 08:52:54

标签: python machine-learning deep-learning pytorch

我有一个预训练的 UNet 模型,其架构如下

def point_function():
    my_list=[]
    num=int(input("Dear User, How many points you want to enter: "))
    for i in range(1,num+1):
        print(f"Please input the coordinates of the point {i}")
        x=float(input("Please enter x-coordinate of the point: "))
        y=float(input("Please enter y-coordinate of the point: "))
        my_tuple=(x,y)
        my_list.append(my_tuple)
    print(my_list)                    

point_function()

该模型采用已使用 min-max 归一化进行归一化的输入。相反,我想在开头添加一个批处理/层规范层,以便我可以在没有规范化的情况下提供图像。

我不想使用 torchvision.transforms 来规范化图像,而是想在开头添加一个对我做同样工作的图层。

1 个答案:

答案 0 :(得分:0)

您可以使用 nn.Module 包装您的预训练模型,该 UNet 将在其前向定义中使用 class UNetWrapper(nn.Module): def __init__(self, unet): super(UNetWrapper, self).__init__() self.norm = nn.BatchNorm2d(3) self.unet = unet def forward(self, x): x = self.norm(x) return self.unet(x)

nn.BatchNorm2d

我不知道你想如何从那里开始,要么微调第一层(可能没那么容易......)。或者自己设置 .eval 参数并在整个包装器上应用 UNet 以锁定批规范层和 {{1}} 以进行评估。