Передача предварительно обученных весов в CNN Pytorch в CNN в Tensorflow

Я обучил эту сеть в Pytorch для изображений размером 224x224 и 4 классов.

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])

Я сохранил это как:

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

Это моя модель, которую я построил в Tensorflow. Я хотел бы передать те веса, которые я сохранил из модели Pytorch, в этот 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? Спасибо!


person Liz    schedule 21.04.2020    source источник


Ответы (1)


Вы можете довольно просто проверить формы всех толщин всех keras слоев:

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

Это даст вам формы всех весов (включая смещения), так что вы можете соответственно подготовить загруженные numpy веса.

Чтобы установить их, сделайте что-то подобное:

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

где torch_weights должен быть списком, содержащим списки np.array, которые вам нужно будет загрузить.

Обычно каждый элемент torch_weights содержит один np.array для весов и один для смещения.

Помните, что формы, полученные при печати, должны быть точно такими же, как те, которые вы добавили set_weights.

Дополнительную информацию см. В документации.

КСТАТИ. Точные формы зависят от слоев и операций, выполняемых моделью, иногда вам может потребоваться транспонировать некоторые массивы, чтобы «подогнать их под них».

person Szymon Maszke    schedule 21.04.2020
comment
Спасибо! У меня уже есть правильные формы для весов в Tensorflow (3,3,3,54). Теперь мой вопрос будет заключаться в том, как передать их моему первому Conv2D слою в Tensorflow. Любые идеи? - person Liz; 22.04.2020
comment
@Liz, как указано, в данном случае это будет model.layers[0].set_weights([your_tensor_with_correct_shape]). Имейте в виду, что torch имеет веса в формате [54, 3, 3, 3] в вашем случае, поскольку это каналы в первую очередь (как минимум до 1.5), поэтому вам необходимо транспонировать его перед сохранением как np.array. weight_tensor.permute(1, 2, 3, 0). - person Szymon Maszke; 22.04.2020
comment
Спасибо!! Теперь он выдает следующую ошибку: ValueError: You called `set_weights(weights)` on layer "conv2d_3" with a weight list of length 3, but the layer was expecting 2 weights. Provided weights: [[[[-0.15836713 -0.178757 0.16782044 ... 0.175... Есть какие-нибудь подсказки, в чем может быть проблема? - person Liz; 22.04.2020
comment
Слишком много аргументов для set_weights, как описано, проверьте, какие веса вы пытаетесь указать - person Szymon Maszke; 22.04.2020