将CNN Pytorch中的预训练砝码传递到Tensorflow中的CNN [英] Pass pretrained weights in CNN Pytorch to a CNN in Tensorflow

问题描述

我已经在Pytorch中为该网络训练了224x224尺寸的图像和4个类.

I have trained this network in Pytorch for 224x224 size images and 4 classes.

class CustomConvNet(nn.Module):
    def __init__(self, num_classes):
        super(CustomConvNet, self).__init__()

        self.layer1 = self.conv_module(3, 64)
        self.layer2 = self.conv_module(64, 128)
        self.layer3 = self.conv_module(128, 256)
        self.layer4 = self.conv_module(256, 256)
        self.layer5 = self.conv_module(256, 512)
        self.gap = self.global_avg_pool(512, num_classes)
        #self.linear = nn.Linear(512, num_classes)
        #self.relu = nn.ReLU()
        #self.softmax = nn.Softmax()

    def forward(self, x):
        out = self.layer1(x)
        out = self.layer2(out)
        out = self.layer3(out)
        out = self.layer4(out)
        out = self.layer5(out)
        out = self.gap(out)
        out = out.view(-1, 4)
        #out = self.linear(out)

        return out

    def conv_module(self, in_num, out_num):
        return nn.Sequential(
            nn.Conv2d(in_num, out_num, kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=(2, 2), stride=None))

    def global_avg_pool(self, in_num, out_num):
        return nn.Sequential(
            nn.Conv2d(in_num, out_num, kernel_size=3, stride=1, padding=1),
            #nn.BatchNorm2d(out_num),
            #nn.LeakyReLU(),

            nn.ReLU(),
            nn.Softmax(),
            nn.AdaptiveAvgPool2d((1, 1)))

我从第一个Conv2D中获得了权重，它的大小为 torch.Size([64，3，3，3])

I got the weights from the first Conv2D and it's size torch.Size([64, 3, 3, 3])

我将其另存为:

weightsCNN = net.layer1[0].weight.data
np.save('CNNweights.npy', weightsCNN)

这是我在Tensorflow中构建的模型.我想将从Pytorch模型中保存的权重传递到此Tensorflow CNN中.

This is my model I built in Tensorflow. I would like to pass those weights I saved from the Pytorch model into this Tensorflow CNN.

    model = models.Sequential()
    model.add(layers.Conv2D(64, (3, 3), activation='relu', input_shape=(224, 224, 3)))
    model.add(layers.MaxPooling2D((2, 2)))

    model.add(layers.Conv2D(128, (3, 3), activation='relu'))
    model.add(layers.MaxPooling2D((2, 2)))

    model.add(layers.Conv2D(256, (3, 3), activation='relu'))
    model.add(layers.MaxPooling2D((2, 2)))

    model.add(layers.Conv2D(256, (3, 3), activation='relu'))
    model.add(layers.MaxPooling2D((2, 2)))

    model.add(layers.Conv2D(512, (3, 3), activation='relu'))
    model.add(layers.MaxPooling2D((2, 2)))

    model.add(layers.Conv2D(512, (3, 3), activation='relu'))

    model.add(layers.GlobalAveragePooling2D())
    model.add(layers.Dense(4, activation='softmax'))
    print(model.summary())


    adam = optimizers.Adam(learning_rate=0.0001, amsgrad=False)
    model.compile(loss='categorical_crossentropy',
                  optimizer=adam,
                  metrics=['accuracy'])


    nb_train_samples = 6596
    nb_validation_samples = 1290
    epochs = 10
    batch_size = 256


    history = model.fit_generator(
        train_generator,
        steps_per_epoch=np.ceil(nb_train_samples/batch_size),
        epochs=epochs,
        validation_data=validation_generator,
        validation_steps=np.ceil(nb_validation_samples / batch_size)
        )

实际上我该怎么做?Tensorflow需要什么形状的砝码?谢谢！

How should I actually do that? What shape of weights does Tensorflow require? Thanks!

推荐答案

您可以非常简单地检查所有 keras 层的所有权重的形状:

You can check shapes of all weights of all keras layers quite simply:

for layer in model.layers:
    print([tensor.shape for tensor in layer.get_weights()])

这将为您提供所有权重(包括偏差)的形状，因此您可以相应地准备加载的 numpy 权重.

This would give you shapes of all weights (including biases), so you can prepare loaded numpy weights accordingly.

要设置它们，请执行类似的操作:

To set them, do something similar:

for torch_weight, layer in zip(model.layers, torch_weights):
    layer.set_weights(torch_weight)

其中 torch_weights 应该是包含要加载的 np.array 列表的列表.

where torch_weights should be a list containing lists of np.array which you would have to load.

通常，每个 torch_weights 的元素将包含一个 np.array 用于权重，一个用于偏置.

Usually each element of torch_weights would contain one np.array for weights and one for bias.

记住从打印中收到的形状必须与您在 set_weights 中放入的形状完全相同.

Remember shapes received from print have to be exactly the same as the ones you put in set_weights.

有关更多信息，请参见文档.

See documentation for more info.

顺便说一句.确切的形状取决于图层和模型执行的操作，有时可能需要转置一些数组以适合它们".

BTW. Exact shapes are dependent on layers and operations performed by model, you may have to transpose some arrays sometimes to "fit them in".

这篇关于将CNN Pytorch中的预训练砝码传递到Tensorflow中的CNN的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！