Как правильно восстановить модель тензорного потока после изменения выходного размера (количества классов)?

Question

Я пытаюсь использовать некоторую существующую модель тензорного потока для решения моей проблемы классификации в некоторых моих собственных наборах данных, вот мой код

import os
from nets import inception

image_size = inception.inception_v1.default_image_size

def get_init_fn():
    """Returns a function run by the chief worker to warm-start the training."""
    checkpoint_exclude_scopes=["InceptionV1/Logits", "InceptionV1/AuxLogits"]

    exclusions = [scope.strip() for scope in checkpoint_exclude_scopes]

    variables_to_restore = []
    for var in slim.get_model_variables():
        excluded = False
        for exclusion in exclusions:
            if var.op.name.startswith(exclusion):
                excluded = True
                break
        if not excluded:
            variables_to_restore.append(var)

    for var in variables_to_restore:
        print(var)

    return slim.assign_from_checkpoint_fn(
      os.path.join(checkpoints_dir, 'inception_v1.ckpt'),
      variables_to_restore)


train_dir = '/tmp/inception_finetuned/'

with tf.Graph().as_default():
    tf.logging.set_verbosity(tf.logging.INFO)

    dataset = custom.get_split('train', custom_data_dir)
    images, _, labels = load_batch(dataset, height=image_size, width=image_size)

    # Create the model, use the default arg scope to configure the batch norm parameters.
    with slim.arg_scope(inception.inception_v1_arg_scope()):
        logits, _ = inception.inception_v1(images, num_classes=dataset.num_classes, is_training=True)

    print('num_classes %d' % dataset.num_classes)
    # Specify the loss function:
    one_hot_labels = slim.one_hot_encoding(labels, dataset.num_classes)
    slim.losses.softmax_cross_entropy(logits, one_hot_labels)
    total_loss = slim.losses.get_total_loss()


    # Create some summaries to visualize the training process:
    tf.summary.scalar('losses/Total_Loss', total_loss)


    # Specify the optimizer and create the train op:
    optimizer = tf.train.AdamOptimizer(learning_rate=0.01)
    train_op = slim.learning.create_train_op(total_loss, optimizer)

    # Run the training:
    final_loss = slim.learning.train(
        train_op,
        logdir=train_dir,
        init_fn=get_init_fn(),
        number_of_steps=2)

print('Finished training. Last batch loss %f' % final_loss)

В начале я попробовал этот код на наборе данных цветов, и он работает как шарм. Затем я попробовал этот код в моем пользовательском наборе данных, а затем внезапно получил эту ошибку:

ИНФОРМАЦИЯ: тензор потока: восстановление параметров из /tmp/inception_finetuned/model.ckpt-2 ИНФОРМАЦИЯ: тензор потока: ошибка сообщена координатору:, При назначении требуются формы обоих тензоров для сопоставления. lhs shape= [1,1,1024,5270] rhs shape= [1,1,1024,5] [[Узел: save_1/Assign_28 = Назначить [T=DT_FLOAT, _class=["loc:@InceptionV1/Logits/Conv2d_0c_1x1/weights"], use_locking=true, validate_shape=true, _device="/job:localhost/replica:0/task:0/gpu:0"](InceptionV1/Logits/Conv2d_0c_1x1/weights/Adam, save_1/RestoreV2_28/_7)]]

Кажется очевидным, что тензор потока обнаружил некоторую контрольную точку во время работы, поэтому он попытался загрузить контрольную точку, но некоторые размеры переменных не соответствуют моей модели с пользовательским размером вывода. Но так как у меня есть get_init_fun(), специально используемый для пропуска несовместимых переменных, я не знаю, что может вызвать ошибку.

Ребята, можете ли вы пролить свет на этот вопрос?