TensorFlow训练自己数据的分类模型

代码连接:
教程视频:

一、数据集的准备和处理

整个项目的目录如下:

|-01_classification_network
    |--data
       |---train
           |---- XX_X.jpg
           ...
       |---test
           |---- XX_X.jpg
           ...
    |--logs
    |--input_data.py
    |--model.py
    |--test_all_image.py
    |--test_one_image.py
    |--training.py

在这里插入图片描述

1、数据集的准备

首先在 data 文件夹里创建一个 train 文件夹, train 文件夹里存放要训练的数据集,数据集的命名方式用的是 “类别_排序.jpg的方式”,也可以按照你自己的要求来命名。
在这里插入图片描述

2、创建加载数据的程序,并测试

data 文件夹的同目录下创建 input_data.py 文件。
在这里插入图片描述
1、加载图片
2、获取batch
3、测试下获取batch成功没(测试完注释掉)

import tensorflow as tf
import numpy as np
import os
import matplotlib.pyplot as plt

#设置训练集所在的路径
train_dir = 'C:/Users/huan4/Desktop/Python/TensorFlow/train_ourself_data/01_classification_network/data/train/'

def get_files(file_dir):
    '''
    Args:
        file_dir: file directory
    Returns:
        list of images and labels
    '''
    #设置每个类别的图片和标签数组
    airplane = []
    label_airplane = []
    car = []
    label_car = []
    bird = []
    label_bird = []
    for file in os.listdir(file_dir):
        name = file.split(sep='_')
        if name[0]=='airplane':
            airplane.append(file_dir + file)
            label_airplane.append(0)
        elif name[0] == 'car':
            car.append(file_dir + file)
            label_car.append(1)
        elif name[0] == 'bird':
            bird.append(file_dir + file)
            label_bird.append(2)

    print('There are %d airplane\nThere are %d car\nThere are %d bird' %(len(airplane), len(car), len(bird))) #输出每个种类图片的数量
    
    image_list = np.hstack((airplane, car, bird))   #np.hstack 水平拼接这些数组
    label_list = np.hstack((label_airplane, label_car, label_bird))

    #使用numpy打乱图片和标签的顺序
    temp = np.array([image_list, label_list])
    temp = temp.transpose()
    np.random.shuffle(temp)

    #再分别把图片和标签取出来,同时把标签int()化
    image_list = list(temp[:, 0])
    label_list = list(temp[:, 1])
    label_list = [int(i) for i in label_list]
    
    return image_list, label_list


def get_batch(image, label, image_W, image_H, batch_size, capacity):
    '''
    Args:
        image: list type
        label: list type
        image_W: image width
        image_H: image height
        batch_size: batch size
        capacity: the maximum elements in queue
    Returns:
        image_batch: 4D tensor [batch_size, width, height, 3], dtype=tf.float32
        label_batch: 1D tensor [batch_size], dtype=tf.int32
    '''

    #将numpy的数据格式转成 tf 的数据格式
    image = tf.cast(image, tf.string)
    label = tf.cast(label, tf.int32)

    # 生成输入的队列
    input_queue = tf.train.slice_input_producer([image, label])
    
    label = input_queue[1]
    image_contents = tf.read_file(input_queue[0])
    image = tf.image.decode_jpeg(image_contents, channels=3)  #将队列里的图片解码
    
    ######################################
    # data argumentation should go to here 数据增强的可以加在这里
    ######################################

    #将图片resize为相同大小
    image = tf.image.resize_image_with_crop_or_pad(image, image_W, image_H)

    #图片归一化
    image = tf.image.per_image_standardization(image)
    
    image_batch, label_batch = tf.train.batch([image, label],
                                                batch_size= batch_size,
                                                num_threads= 64, 
                                                capacity = capacity)
    
    #you can also use shuffle_batch  打乱数据的顺序,因为前面打乱了,这里就不打乱了
#    image_batch, label_batch = tf.train.shuffle_batch([image,label],
#                                                      batch_size=BATCH_SIZE,
#                                                      num_threads=64,
#                                                      capacity=CAPACITY,
#                                                      min_after_dequeue=CAPACITY-1)
    
    label_batch = tf.reshape(label_batch, [batch_size])
    image_batch = tf.cast(image_batch, tf.float32)
    
    return image_batch, label_batch


def test_batch():

    BATCH_SIZE = 2
    CAPACITY = 256
    IMG_W = 32
    IMG_H = 32

    image_list, label_list = get_files(train_dir)
    image_batch, label_batch = get_batch(image_list, label_list, IMG_W, IMG_H, BATCH_SIZE, CAPACITY)

    with tf.Session() as sess:
       i = 0   #只看看第一个batch就好
       coord = tf.train.Coordinator()          #读取batch的操作,入列和出列
       threads = tf.train.start_queue_runners(coord=coord)

       try:
           while not coord.should_stop() and i<1:

               img, label = sess.run([image_batch, label_batch])

               # just test one batch
               for j in np.arange(BATCH_SIZE):
                   print('label: %d' %label[j])
                   plt.imshow(img[j,:,:,:])
                   plt.show()
               i+=1

       except tf.errors.OutOfRangeError:
           print('done!')
       finally:
           coord.request_stop()
       coord.join(threads)

#test_batch()  #测试的时候用,训练的时候注释掉

测试效果:
在这里插入图片描述
这是归一化的图片,想看原图的话在get_batch()函数里把图片归一化注释掉image = tf.image.per_image_standardization(image),但是训练的时候必须用归一化,这里注释只是为了显示原图。

原图:
在这里插入图片描述

二、分类模型的建立

data 文件夹的同目录下创建 model.py 文件。
在这里插入图片描述
1、 定义网络的结构
conv1 - pool1 - conv2 - pool2 - fc3 - fc4 - softmax
2、 计算 loss
3、 训练,设置学习率,使 loss 降低
4、 评估准确率

import tensorflow as tf

def inference(images, batch_size, n_classes):
    '''Build the model
    Args:
        images: image batch, 4D tensor, tf.float32, [batch_size, width, height, channels]
    Returns:
        output tensor with the computed logits, float, [batch_size, n_classes]
    '''
    
    with tf.variable_scope('conv1') as scope:
        weights = tf.get_variable('weights', 
                                  shape = [3,3,3, 16],    #shape = [kernel size, kernel size, channels(图片的通道数), kernel numbers(卷积核的数量)]
                                  dtype = tf.float32, 
                                  initializer=tf.truncated_normal_initializer(stddev=0.1,dtype=tf.float32))   #初始化
        biases = tf.get_variable('biases', 
                                 shape=[16],  #偏置值,和这层的卷积核数量相同
                                 dtype=tf.float32,
                                 initializer=tf.constant_initializer(0.1))  #初始化
        conv = tf.nn.conv2d(images, weights, strides=[1,1,1,1], padding='SAME')
        pre_activation = tf.nn.bias_add(conv, biases)
        conv1 = tf.nn.relu(pre_activation, name= scope.name)
    
    #pool1 and norm1   
    with tf.variable_scope('pooling1_lrn') as scope:
        pool1 = tf.nn.max_pool(conv1, ksize=[1,3,3,1],strides=[1,2,2,1],
                               padding='SAME', name='pooling1')
        norm1 = tf.nn.lrn(pool1, depth_radius=4, bias=1.0, alpha=0.001/9.0,
                          beta=0.75,name='norm1')
    
    #conv2
    with tf.variable_scope('conv2') as scope:
        weights = tf.get_variable('weights',
                                  shape=[3,3,16,16],
                                  dtype=tf.float32,
                                  initializer=tf.truncated_normal_initializer(stddev=0.1,dtype=tf.float32))
        biases = tf.get_variable('biases',
                                 shape=[16], 
                                 dtype=tf.float32,
                                 initializer=tf.constant_initializer(0.1))
        conv = tf.nn.conv2d(norm1, weights, strides=[1,1,1,1],padding='SAME')
        pre_activation = tf.nn.bias_add(conv, biases)
        conv2 = tf.nn.relu(pre_activation, name='conv2')
    
    
    #pool2 and norm2
    with tf.variable_scope('pooling2_lrn') as scope:
        norm2 = tf.nn.lrn(conv2, depth_radius=4, bias=1.0, alpha=0.001/9.0,
                          beta=0.75,name='norm2')
        pool2 = tf.nn.max_pool(norm2, ksize=[1,3,3,1], strides=[1,1,1,1],
                               padding='SAME',name='pooling2')
    
    
    #local3
    with tf.variable_scope('local3') as scope:
        reshape = tf.reshape(pool2, shape=[batch_size, -1])
        dim = reshape.get_shape()[1].value
        weights = tf.get_variable('weights',
                                  shape=[dim,128],
                                  dtype=tf.float32,
                                  initializer=tf.truncated_normal_initializer(stddev=0.005,dtype=tf.float32))
        biases = tf.get_variable('biases',
                                 shape=[128],
                                 dtype=tf.float32, 
                                 initializer=tf.constant_initializer(0.1))
        local3 = tf.nn.relu(tf.matmul(reshape, weights) + biases, name=scope.name)    
    
    #local4
    with tf.variable_scope('local4') as scope:
        weights = tf.get_variable('weights',
                                  shape=[128,128],
                                  dtype=tf.float32, 
                                  initializer=tf.truncated_normal_initializer(stddev=0.005,dtype=tf.float32))
        biases = tf.get_variable('biases',
                                 shape=[128],
                                 dtype=tf.float32,
                                 initializer=tf.constant_initializer(0.1))
        local4 = tf.nn.relu(tf.matmul(local3, weights) + biases, name='local4')
     
        
    # softmax
    with tf.variable_scope('softmax_linear') as scope:
        weights = tf.get_variable('softmax_linear',
                                  shape=[128, n_classes],
                                  dtype=tf.float32,
                                  initializer=tf.truncated_normal_initializer(stddev=0.005,dtype=tf.float32))
        biases = tf.get_variable('biases', 
                                 shape=[n_classes],
                                 dtype=tf.float32, 
                                 initializer=tf.constant_initializer(0.1))
        softmax_linear = tf.add(tf.matmul(local4, weights), biases, name='softmax_linear')
    
    return softmax_linear

#%%
def losses(logits, labels):
    '''Compute loss from logits and labels
    Args:
        logits: logits tensor, float, [batch_size, n_classes]
        labels: label tensor, tf.int32, [batch_size]
        
    Returns:
        loss tensor of float type
    '''
    with tf.variable_scope('loss') as scope:
        cross_entropy = tf.nn.sparse_softmax_cross_entropy_with_logits\
                        (logits=logits, labels=labels, name='xentropy_per_example')
        loss = tf.reduce_mean(cross_entropy, name='loss')
        tf.summary.scalar(scope.name+'/loss', loss)
    return loss


def trainning(loss, learning_rate):
    '''Training ops, the Op returned by this function is what must be passed to 
        'sess.run()' call to cause the model to train.
        
    Args:
        loss: loss tensor, from losses()
        
    Returns:
        train_op: The op for trainning
    '''
    with tf.name_scope('optimizer'):
        optimizer = tf.train.AdamOptimizer(learning_rate= learning_rate)
        global_step = tf.Variable(0, name='global_step', trainable=False)
        train_op = optimizer.minimize(loss, global_step= global_step)
    return train_op


def evaluation(logits, labels):
  """Evaluate the quality of the logits at predicting the label.
  Args:
    logits: Logits tensor, float - [batch_size, NUM_CLASSES].
    labels: Labels tensor, int32 - [batch_size], with values in the
      range [0, NUM_CLASSES).
  Returns:
    A scalar int32 tensor with the number of examples (out of batch_size)
    that were predicted correctly.
  """
  with tf.variable_scope('accuracy') as scope:
      correct = tf.nn.in_top_k(logits, labels, 1)
      correct = tf.cast(correct, tf.float16)
      accuracy = tf.reduce_mean(correct)
      tf.summary.scalar(scope.name+'/accuracy', accuracy)
  return accuracy

三、训练模型

data 文件夹的同目录下创建 training.py 文件。
在这里插入图片描述

import os
import numpy as np
import tensorflow as tf
import input_data
import model


N_CLASSES = 3  #设置你自己数据集的类别数
IMG_W = 32
IMG_H = 32
BATCH_SIZE = 16
CAPACITY = 2000
MAX_STEP = 10000 # 设置训练的步数,每步训练一个 batch_size
learning_rate = 0.0001 # 学习率的设置


def run_training():

    train_dir = 'C:/Users/huan4/Desktop/Python/TensorFlow/train_ourself_data/01_classification_network/data/train/'
    logs_train_dir ='C:/Users/huan4/Desktop/Python/TensorFlow/train_ourself_data/01_classification_network/logs/train/'
    
    train, train_label = input_data.get_files(train_dir)
    
    train_batch, train_label_batch = input_data.get_batch(train,
                                                          train_label,
                                                          IMG_W,
                                                          IMG_H,
                                                          BATCH_SIZE, 
                                                          CAPACITY)      
    train_logits = model.inference(train_batch, BATCH_SIZE, N_CLASSES)
    train_loss = model.losses(train_logits, train_label_batch)        
    train_op = model.trainning(train_loss, learning_rate)
    train__acc = model.evaluation(train_logits, train_label_batch)
       
    summary_op = tf.summary.merge_all()
    sess = tf.Session()
    train_writer = tf.summary.FileWriter(logs_train_dir, sess.graph)
    saver = tf.train.Saver()
    
    sess.run(tf.global_variables_initializer())
    coord = tf.train.Coordinator()
    threads = tf.train.start_queue_runners(sess=sess, coord=coord)
    
    try:
        for step in np.arange(MAX_STEP):
            if coord.should_stop():
                    break

            _, tra_loss, tra_acc = sess.run([train_op, train_loss, train__acc])
               
            if step % 50 == 0:
                print('Step %d, train loss = %.2f, train accuracy = %.2f%%' %(step, tra_loss, tra_acc*100.0))
                summary_str = sess.run(summary_op)
                train_writer.add_summary(summary_str, step)
            
            if step % 2000 == 0 or (step + 1) == MAX_STEP:
                checkpoint_path = os.path.join(logs_train_dir, 'model.ckpt')
                saver.save(sess, checkpoint_path, global_step=step)
                
    except tf.errors.OutOfRangeError:
        print('Done training -- epoch limit reached')
    finally:
        coord.request_stop()
        
    coord.join(threads)
    sess.close()

run_training()

训练记录
在这里插入图片描述

四、训练测试

4.1、测试单张图片

data 文件夹的同目录下创建 test_one_image.py 文件。
在这里插入图片描述

import numpy as np
import tensorflow as tf
import input_data
import model
from PIL import Image
import matplotlib.pyplot as plt

def get_one_image(train, train_label):
   '''Randomly pick one image from training data
   Return: ndarray
   '''
   n = len(train)
   ind = np.random.randint(0, n)
   img_dir = train[ind]
   label = train_label[ind]
   image = Image.open(img_dir)
   plt.imshow(image)
   image = image.resize([32, 32])
   image = np.array(image)
   return image,label,img_dir

def evaluate_one_image():

   train_dir = 'C:/Users/huan4/Desktop/Python/TensorFlow/train_ourself_data/01_classification_network/data/train/'
   train, train_label = input_data.get_files(train_dir)
   image_array, label, img_dir = get_one_image(train, train_label)
   name = img_dir.split(sep='/')
   with tf.Graph().as_default():
       BATCH_SIZE = 1
       N_CLASSES = 3

       image = tf.cast(image_array, tf.float32)
       image = tf.image.per_image_standardization(image)
       image = tf.reshape(image, [1, 32, 32, 3])
       logit = model.inference(image, BATCH_SIZE, N_CLASSES)

       logit = tf.nn.softmax(logit)

       x = tf.placeholder(tf.float32, shape=[32, 32, 3])

       logs_train_dir = 'C:/Users/huan4/Desktop/Python/TensorFlow/train_ourself_data/01_classification_network/logs/train/'

       saver = tf.train.Saver()

       with tf.Session() as sess:

           print("Reading checkpoints...")
           ckpt = tf.train.get_checkpoint_state(logs_train_dir)
           if ckpt and ckpt.model_checkpoint_path:
               global_step = ckpt.model_checkpoint_path.split('/')[-1].split('-')[-1]
               saver.restore(sess, ckpt.model_checkpoint_path)
               print('Loading success, global_step is %s' % global_step)
           else:
               print('No checkpoint file found')

           prediction = sess.run(logit, feed_dict={x: image_array})
           max_index = np.argmax(prediction)

           print('This image is %s, lable is %s' % (name[-1], label))
           if max_index==0:
               print('This is a airplane with possibility %.6f' %prediction[:, 0])
           elif max_index==1:
               print('This is a car with possibility %.6f' %prediction[:, 1])
           elif max_index==2:
               print('This is a bird with possibility %.6f' % prediction[:, 2])

evaluate_one_image()

测试结果:
在这里插入图片描述

4.2、测试所有图片

./data/test 文件夹里存放所有要测试的图片,注意测试图片的命名方式和训练的图片命名方式相同。
data 文件夹的同目录下创建 test_all_image.py 文件。
在这里插入图片描述
在这里插入图片描述
代码:

import tensorflow as tf
import input_data
import model

IMG_W = 32
IMG_H = 32
CAPACITY = 2000

def evaluate_all_image():

    test_dir = 'C:/Users/huan4/Desktop/Python/TensorFlow/train_ourself_data/01_classification_network/data/test/'
    N_CLASSES = 3
    print('-------------------------')
    test, test_label = input_data.get_files(test_dir)
    BATCH_SIZE = len(test)
    print('There are %d test images totally..' % BATCH_SIZE)
    print('-------------------------')
    test_batch, test_label_batch = input_data.get_batch(test,
                                                          test_label,
                                                          IMG_W,
                                                          IMG_H,
                                                          BATCH_SIZE, 
                                                          CAPACITY)
	
    logits = model.inference(test_batch, BATCH_SIZE, N_CLASSES)
    testloss = model.losses(logits, test_label_batch) 
    testacc = model.evaluation(logits, test_label_batch)
	
    logs_train_dir = 'C:/Users/huan4/Desktop/Python/TensorFlow/train_ourself_data/01_classification_network/logs/train/'
    saver = tf.train.Saver()
        
    with tf.Session() as sess:
        print("Reading checkpoints...")
        ckpt = tf.train.get_checkpoint_state(logs_train_dir)
        if ckpt and ckpt.model_checkpoint_path:
            global_step = ckpt.model_checkpoint_path.split('/')[-1].split('-')[-1]
            saver.restore(sess, ckpt.model_checkpoint_path)
            print('Loading success, global_step is %s' % global_step)
        else:
            print('No checkpoint file found')
        print('-------------------------')        
        coord = tf.train.Coordinator()
        threads = tf.train.start_queue_runners(sess=sess, coord=coord)
        test_loss,test_acc = sess.run([testloss,testacc])
        print('The model\'s loss is %.2f' %test_loss)
        correct = int(BATCH_SIZE*test_acc)
        print('Correct : %d' % correct)
        print('Wrong : %d' % (BATCH_SIZE - correct))
        print('The accuracy in test images are %.2f%%' %(test_acc*100.0))
    coord.request_stop()
    coord.join(threads)
    sess.close()

evaluate_all_image()

测试结果:
在这里插入图片描述

  • 8
    点赞
  • 100
    收藏
    觉得还不错? 一键收藏
  • 7
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论 7
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值