Skip to main content
 首页 » 编程设计

python之Tensorflow Slim 中的批量训练

2025年02月15日28zlslch

我正在查看 TF Slim介绍性文件,据我所知,它每次运行只接收一批图像数据(32 张图像)。显然,人们想要遍历它并训练许多不同的批处理。简介不包括这一点。如何正确完成此操作。我想应该有一些方法可以指定一个加载批处理函数,该函数应该在开始批处理训练事件时自动调用,但我似乎无法在介绍中找到一个简单的例子。

# Note that this may take several minutes. 
 
import os 
 
from datasets import flowers 
from nets import inception 
from preprocessing import inception_preprocessing 
 
slim = tf.contrib.slim 
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) 
 
    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 = flowers.get_split('train', flowers_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) 
 
    # 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.scalar_summary('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) 

请您参考如下方法:

slim.learning.train 函数包含一个训练循环,因此您提供的代码实际上是针对多批图像进行训练的。

参见 here in the source code ,其中 train_step_fn 在 while 循环中被调用。 train_step(train_step_fn 的默认值)包含行 sess.run([train_op, global_step]...),它实际运行对单批图像的训练操作。