tensorflow、keras和pytorch搭建DNN、CNN、RNN手写数字识别
MNIST手写数字集
MNIST是一个由美国由美国邮政系统开发的手写数字识别数据集。手写内容是0~9，一共有60000个图片样本，我们可以到MNIST官网免费下载，总共4个.gz后缀的压缩文件，该文件是二进制内容。
| 文件名 | 大小 | 用途 |
| train-images-idx3-ubyte.gz | 9.45MB | 训练图像数据 |
| train-labels-idx1-ubyte.gz | 0.03MB | 训练图像的标签 |
| t10k-images-idx3-ubyte.gz | 1.57MB | 测试图像数据 |
| t10k-labels-idx1-ubyte.gz | 4.4KB | 测试图像的标签 |
下载MNIST数据集
方法一、官网下载（4个gz文件，图像的取值在0~1之间）
方法二、谷歌下载（1个npz文件，图像的取值在0~255之间）
方法三、通过tensorflow或keras代码获取
from tensorflow.examples.tutorials.mnist import input_data # tensorflow（1.7版本以前） # 从MNIST_data/中读取MNIST数据。当数据不存在时，会自动执行下载 mnist = input_data.read_data_sets("./mnist/", one_hot=True) # tensorflow（1.7版本以后） import tensorflow as tf (train_x, train_y), (test_x, test_y) = tf.keras.datasets.mnist.load_data(path='mnist.npz') # keras代码获取 from keras.datasets import mnist (train_x, train_y), (test_x, test_y) = mnist.load_data() # 通过numpy代码获取.npz中的数据 f = np.load(path) x_train, y_train = f['x_train'], f['y_train'] x_test, y_test = f['x_test'], f['y_test'] f.close()
如果通过代码下载MNIST的方法，不FQ的话，可能无法顺利下载MNSIT数据集，因此我建议大家还是先手动下载好，再来通过代码导入。
MNIST图像
训练数据集包含 60,000 个样本, 测试数据集包含 10,000 样本。在 MNIST 数据集中的每张图片由 28 x 28(=784) 个像素点构成, 每个像素点用一个灰度值表示。
我们可以通过下面python代码下载MNIST数据集，并窥探一下MNIST数据集的内部数据集的划分，以及手写数字的长相。
import matplotlib.pyplot as plt from tensorflow.examples.tutorials.mnist import input_data
# 从MNIST_data/中读取MNIST数据。当数据不存在时，会自动执行下载 mnist = input_data.read_data_sets('./mnist', one_hot=True) # 将数组张换成图片形式 print(mnist.train.images.shape) # 训练数据图片(55000, 784) print(mnist.train.labels.shape) # 训练数据标签(55000, 10) print(mnist.test.images.shape) # 测试数据图片(10000, 784) print(mnist.test.labels.shape) # 测试数据图片(10000, 10) print(mnist.validation.images.shape) # 验证数据图片(5000, 784) print(mnist.validation.labels.shape) # 验证数据图片(5000, 784) print(mnist.train.labels[1]) # [0. 0. 0. 0. 0. 0. 0. 1. 0. 0.] image = mnist.train.images[1].reshape(28, 28) fig = plt.figure("图片展示") plt.imshow(image,cmap='gray') plt.axis('off') #不显示坐标尺寸 plt.show()
在画出数字的同时，同时取出标签．
from tensorflow.examples.tutorials.mnist import input_data import math import matplotlib.pyplot as plt import numpy as np mnist = input_data.read_data_sets('./mnist', one_hot=True) # 画单张mnist数据集的数字 def drawdigit(position,image, title): plt.subplot(*position) # 星号元组传参 plt.imshow(image, cmap='gray_r') plt.axis('off') plt.title(title) # 取一个batch的数据，然后在一张画布上画batch_size个子图 def batchDraw(batch_size): images, labels = mnist.train.next_batch(batch_size) row_num = math.ceil(batch_size ** 0.5) # 向上取整 column_num = row_num plt.figure(figsize=(row_num, column_num)) # 行．列 for i in range(row_num): for j in range(column_num): index = i * column_num + j if index < batch_size: position = (row_num, column_num, index+1) image = images[index].reshape(28, 28) # 取出列表中最大数的索引 title = 'actual:%d' % (np.argmax(labels[index])) drawdigit(position, image, title) if __name__ == '__main__': batchDraw(16) plt.show()
代码说明：
mnist = input_data.read_data_sets("./mnist/", one_hot=True, reshape=False)
图像是由RGB三个数组组成的，而灰度图只是其中一个数组，而图像是由像素组成，每个像素的值在0~225之间，MNIST数据集中的每个数字都有28*28＝784个像素值．上面的代码如果reshape＝True(默认)，MNIST数据的shape=(?, 784)，如果reshape=False MNIST数据为(?, 28,28,1)．
Keras
DNN网络
from keras.models import Model from keras.layers import Input, Dense, Dropout from keras import regularizers from keras.optimizers import Adam from tensorflow.examples.tutorials.mnist import input_data mnist = input_data.read_data_sets("mnist/", one_hot=True) x_train = mnist.train.images # 训练数据 (55000, 784) y_train = mnist.train.labels # 训练标签 x_test = mnist.test.images y_test = mnist.test.images # DNN网络结构 inputs = Input(shape=(784,)) h1 = Dense(64, activation='relu', kernel_regularizer=regularizers.l2(0.01))(inputs) # 权重矩阵l2正则化 h1 = Dropout(0.2)(h1) h2 = Dense(64, activation='relu', kernel_regularizer=regularizers.l2(0.01))(h1) # 权重矩阵l2正则化 h2 = Dropout(0.2)(h2) h3 = Dense(64, activation='relu', kernel_regularizer=regularizers.l2(0.01))(h2) # 权重矩阵l2正则化 h3 = Dropout(0.2)(h3) outputs = Dense(10, activation='softmax', kernel_regularizer=regularizers.l2(0.01))(h3) # 权重矩阵l2正则化 model = Model(input=inputs, output=outputs) # 编译模型 opt = Adam(lr=0.01, beta_1=0.9, beta_2=0.999, epsilon=1e-08) # epsilon模糊因子 model.compile(optimizer=opt, loss='categorical_crossentropy', metrics=['accuracy']) # 交叉熵损失函数 # 开始训练 model.fit(x=x_train, y=y_train, validation_split=0.1, batch_size=128, epochs=4) model.save('k_DNN.h5')
CNN网络
from keras.models import Model from keras.layers import Input, Conv2D, MaxPooling2D, Reshape, Dense from keras import regularizers from keras.optimizers import Adam from tensorflow.examples.tutorials.mnist import input_data mnist = input_data.read_data_sets("./mnist/", one_hot=True, reshape=False) x_train = mnist.train.images # 训练数据 (55000, 28, 28, 1) y_train = mnist.train.labels # 训练标签 x_test = mnist.test.images y_test = mnist.test.images # 网络结构 input = Input(shape=(28, 28, 1)) h1 = Conv2D(filters=64, kernel_size=(3,3), strides=(1, 1), padding='same', activation='relu')(input) h1 = MaxPooling2D(pool_size=2, strides=2, padding='valid')(h1) h1 = Conv2D(filters=32, kernel_size=(3,3), strides=(1, 1), padding='same', activation='relu')(h1) h1 = MaxPooling2D()(h1) h1 = Conv2D(filters=16, kernel_size=(3,3), strides=(1, 1), padding='same', activation='relu')(h1) h1 = Reshape((16 * 7 * 7,))(h1) # h1.shape (?, 16*7*7) output = Dense(10, activation="softmax", kernel_regularizer=regularizers.l2(0.01))(h1) model = Model(input=input, output=output) model.summary() # 编译模型 opt = Adam(lr=0.01, beta_1=0.9, beta_2=0.999, epsilon=1e-08) model.compile(optimizer=opt, loss="categorical_crossentropy", metrics=["accuracy"]) # 开始训练 model.fit(x=x_train, y=y_train, validation_split=0.1, epochs=5) model.save('k_CNN.h5')
RNN网络
from keras.models import Model from keras.layers import Input, LSTM, Dense from keras import regularizers from keras.optimizers import Adam from tensorflow.examples.tutorials.mnist import input_data mnist = input_data.read_data_sets("./mnist/", one_hot=True) x_train = mnist.train.images # (28, 28, 1) x_train = x_train.reshape(-1, 28, 28) y_train = mnist.train.labels # RNN网络结构 inputs = Input(shape=(28, 28)) h1 = LSTM(64, activation='relu', return_sequences=True, dropout=0.2)(inputs) h2 = LSTM(64, activation='relu', dropout=0.2)(h1) outputs = Dense(10, activation='softmax', kernel_regularizer=regularizers.l2(0.01))(h2) model = Model(input=inputs, output=outputs) # 编译模型 opt = Adam(lr=0.003, beta_1=0.9, beta_2=0.999, epsilon=1e-08) model.compile(optimizer=opt, loss='categorical_crossentropy', metrics=['accuracy']) # 训练模型 model.fit(x=x_train, y=y_train, validation_split=0.1, batch_size=128, epochs=5) model.save('k_RNN.h5')
Tensorflow
DNN网络
import tensorflow as tf from tensorflow.examples.tutorials.mnist import input_data mnist = input_data.read_data_sets("./mnist", one_hot=True) # train image shape: (55000, 784) # trian label shape: (55000, 10) # val image shape: (5000, 784) # test image shape: (10000, 784) epochs = 2 output_size = 10 input_size = 784 hidden1_size = 512 hidden2_size = 256 batch_size = 1000 learning_rate_base = 0.005 unit_list = [784, 512, 256, 10] batch_num = mnist.train.labels.shape[0] // batch_size # 全连接神经网络 def dense(x, w, b, keeppord): linear = tf.matmul(x, w) + b activation = tf.nn.relu(linear) y = tf.nn.dropout(activation,keeppord) return y def DNNModel(image, w, b, keeppord): dense1 = dense(image, w[0], b[0],keeppord) dense2 = dense(dense1, w[1], b[1],keeppord) output = tf.matmul(dense2, w[2]) + b[2] return output # 生成网络的权重 def gen_weights(unit_list): w = [] b = [] # 遍历层数 for i in range(len(unit_list)-1): sub_w = tf.Variable(tf.random_normal(shape=[unit_list[i], unit_list[i+1]])) sub_b = tf.Variable(tf.random_normal(shape=[unit_list[i+1]])) w.append(sub_w) b.append(sub_b) return w, b x = tf.placeholder(tf.float32, [None, 784]) y_true = tf.placeholder(tf.float32, [None, 10]) keepprob = tf.placeholder(tf.float32) global_step = tf.Variable(0) w, b = gen_weights(unit_list) y_pre = DNNModel(x, w, b, keepprob) loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits_v2(logits=y_pre, labels=y_true)) tf.summary.scalar("loss", loss) # 收集标量 opt = tf.train.AdamOptimizer(0.001).minimize(loss, global_step=global_step) predict = tf.equal(tf.argmax(y_pre, axis=1), tf.argmax(y_true, axis=1)) # 返回每行或者每列最大值的索引，判断是否相等 acc = tf.reduce_mean(tf.cast(predict, tf.float32)) tf.summary.scalar("acc", acc) # 收集标量 merged = tf.summary.merge_all() # 和并变量 saver = tf.train.Saver() # 保存和加载模型 init = tf.global_variables_initializer() # 初始化全局变量 with tf.Session() as sess: sess.run(init) writer = tf.summary.FileWriter("./logs/tensorboard", tf.get_default_graph()) # tensorboard 事件文件 for i in range(batch_num * epochs): x_train, y_train = mnist.train.next_batch(batch_size) summary, _ = sess.run([merged, opt], feed_dict={x:x_train, y_true:y_train, keepprob: 0.75}) writer.add_summary(summary, i) # 将每次迭代后的变量写入事件文件 # 评估模型在验证集上的识别率 if i % 50 == 0: feeddict = {x: mnist.validation.images, y_true: mnist.validation.labels, keepprob: 1.} # 验证集 valloss, accuracy = sess.run([loss, acc], feed_dict=feeddict) print(i, 'th batch val loss:', valloss, ', accuracy:', accuracy) saver.save(sess, './checkpoints/tfdnn.ckpt') # 保存模型 print('测试集准确度:', sess.run(acc, feed_dict={x:mnist.test.images, y_true:mnist.test.labels, keepprob:1.})) writer.close()
CNN网络
import tensorflow as tf from tensorflow.examples.tutorials.mnist import input_data epochs = 10 batch_size = 100 mnist = input_data.read_data_sets("mnist/", one_hot=True, reshape=False) batch_nums = mnist.train.labels.shape[0] // batch_size # 卷积结构 def conv2d(x, w, b): # x = (?, 28,28,1) # filter = [filter_height, filter_width, in_channels, out_channels] # data_format = [批次，高度，宽度，通道] # 第一个和第四个必须是1 return tf.nn.conv2d(x, filter=w, strides=[1, 1, 1, 1], padding='SAME') + b def pool(x): return tf.nn.max_pool(x, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME') # 定义网络结构 def cnn_net(x, keepprob): # x = reshape=False (?, 28,28,1) w1 = tf.Variable(tf.random_normal([5, 5, 1, 64])) b1 = tf.Variable(tf.random_normal([64])) w2 = tf.Variable(tf.random_normal([5, 5, 64, 32])) b2 = tf.Variable(tf.random_normal([32])) w3 = tf.Variable(tf.random_normal([7 * 7 * 32, 10])) b3 = tf.Variable(tf.random_normal([10])) hidden1 = pool(conv2d(x, w1, b1)) hidden1 = tf.nn.dropout(hidden1, keepprob) hidden2 = pool(conv2d(hidden1, w2, b2)) hidden2 = tf.reshape(hidden2, [-1, 7 * 7 * 32]) hidden2 = tf.nn.dropout(hidden2, keepprob) output = tf.matmul(hidden2, w3) + b3 return output # 定义所需占位符 x = tf.placeholder(tf.float32, [None, 28, 28, 1]) y_true = tf.placeholder(tf.float32, [None, 10]) keepprob = tf.placeholder(tf.float32) # 在训练模型时，随着训练的逐步降低学习率。该函数返回衰减后的学习率。 global_step = tf.Variable(0) learning_rate = tf.train.exponential_decay(0.01, global_step, 100, 0.96, staircase=True) # 训练所需损失函数 logits = cnn_net(x, keepprob) loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits_v2(logits=logits, labels=y_true)) opt = tf.train.AdamOptimizer(learning_rate).minimize(loss, global_step=global_step) # 定义评估模型 predict = tf.equal(tf.argmax(logits, 1), tf.argmax(y_true, 1)) # 预测值 accuracy = tf.reduce_mean(tf.cast(predict, tf.float32)) # 验证值 init = tf.global_variables_initializer() # 开始训练 with tf.Session() as sess: sess.run(init) for k in range(epochs): for i in range(batch_nums): train_x, train_y = mnist.train.next_batch(batch_size) sess.run(opt, {x: train_x, y_true: train_y, keepprob: 0.75}) # 评估模型在验证集上的识别率 if i % 50 == 0: acc = sess.run(accuracy, {x: mnist.validation.images[:1000], y_true: mnist.validation.labels[:1000], keepprob: 1.}) print(k, 'epochs, ', i, 'iters, ', ', acc :', acc)
RNN网络
import tensorflow as tf from tensorflow.examples.tutorials.mnist import input_data epochs = 10 batch_size = 1000 mnist = input_data.read_data_sets("mnist/", one_hot=True) batch_nums = mnist.train.labels.shape[0] // batch_size # 定义网络结构 def RNN_Model(x, batch_size, keepprob): # rnn_layers = [tf.nn.rnn_cell.LSTMCell(size) for size in [28, 28]] rnn_cell = tf.nn.rnn_cell.LSTMCell(28) rnn_drop = tf.nn.rnn_cell.DropoutWrapper(rnn_cell, output_keep_prob=keepprob) # 创建由多个RNNCell组成的RNN单元。 multi_rnn_cell = tf.nn.rnn_cell.MultiRNNCell([rnn_drop] * 2) initial_state = multi_rnn_cell.zero_state(batch_size, tf.float32) # 创建由RNNCell指定的递归神经网络cell。执行完全动态展开inputs outputs, states = tf.nn.dynamic_rnn(cell=multi_rnn_cell, inputs=x, dtype=tf.float32, initial_state=initial_state ) # outputs 的shape为[batch_size, max_time, 28] w = tf.Variable(tf.random_normal([28, 10])) b = tf.Variable(tf.random_normal([10])) output = tf.matmul(outputs[:, -1, :], w) + b return output, states # 定义所需占位符 x = tf.placeholder(tf.float32, [None, 28, 28]) y_true = tf.placeholder(tf.float32, [None, 10]) keepprob = tf.placeholder(tf.float32) global_step = tf.Variable(0) # 在训练模型时，随着训练的逐步降低学习率。该函数返回衰减后的学习率。 learning_rate = tf.train.exponential_decay(0.01, global_step, 10, 0.96, staircase=True) # 训练所需损失函数 y_pred, states = RNN_Model(x, batch_size, keepprob) loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=y_pred, labels=y_true)) opt = tf.train.AdamOptimizer(learning_rate).minimize(loss, global_step=global_step) # 最小化损失函数 predict = tf.equal(tf.argmax(y_pred, 1), tf.argmax(y_true, 1)) # 预测值 acc = tf.reduce_mean(tf.cast(predict, tf.float32)) # 精度 init = tf.global_variables_initializer() # 开始训练 with tf.Session() as sess: sess.run(init) for k in range(epochs): for i in range(batch_nums): train_x, train_y = mnist.train.next_batch(batch_size) sess.run(opt, {x: train_x.reshape((-1, 28, 28)), y_true: train_y, keepprob: 0.8}) # 评估模型在验证集上的识别率 if i % 50 == 0: val_losses = 0 accuracy = 0 val_x, val_y = mnist.validation.next_batch(batch_size) for i in range(val_x.shape[0]): val_loss, accy = sess.run([loss, acc], {x: val_x.reshape((-1, 28, 28)), y_true: val_y, keepprob: 1.}) val_losses += val_loss accuracy += accy print('val_loss is :', val_losses / val_x.shape[0], ', accuracy is :', accuracy / val_x.shape[0])
pytorch
DNN网络
class MNISTNet(nn.Module): def __init__(self, input_dims, n_hiddens, n_class): super(MNISTNet, self).__init__() self.block_1 = nn.Sequential( nn.Linear(in_features=input_dims, out_features=n_hiddens), nn.ReLU(), nn.Dropout(p=0.2) ) self.block_2 = nn.Sequential( nn.Linear(in_features=n_hiddens, out_features=n_hiddens), nn.ReLU(), nn.Dropout(p=0.2) ) self.last_Linear = nn.Linear(in_features=n_hiddens, out_features=n_class) def forward(self, inputs): x = inputs.view(inputs.size(0), -1) # (batch,28*28) x = self.block_1(x) x = self.block_2(x) x = self.last_Linear(x) return x def test_MNISTNet(): inputs = torch.randn(64, 1, 28, 28) model = MNISTNet(input_dims=784, n_hiddens=256, n_class=10) outputs = model(inputs) print("outputs", outputs.shape)
pytorch模型训练
# -*- coding:utf-8 -*- # Author:凌逆战 | Never # Date: 2023/1/19 """ 参考：https://github.com/aaron-xichen/pytorch-playground """ import os # os.environ["CUDA_VISIBLE_DEVICES"] = "0" # 使用第一个和第二个GPU import argparse import torch from torch.utils.data import DataLoader from tensorboardX import SummaryWriter from model.demo_model import MNISTNet import torchvision import torchvision.transforms as transforms import torch.nn.functional as F def parse_args(): parser = argparse.ArgumentParser(description='PyTorch MNIST Example') parser.add_argument("--model_class", type=str, default="MNISTNet", help="模型类名") parser.add_argument('--train_tag', default="Mask(IAM)_Loss(mask_MAE)", help='训练标记') parser.add_argument("--model_name", type=str, default=None, help="是否加载模型继续训练 '10.pth' None") parser.add_argument("--batch_size", type=int, default=64, help="") parser.add_argument("--epochs", type=int, default=100) parser.add_argument('--lr', type=float, default=3e-4, help='学习率 (default: 0.01)') parser.add_argument('--train_log_dir', default="./train_log", help='训练记录文件夹') args = parser.parse_args() return args def main(): args = parse_args() print("GPU是否可用：", torch.cuda.is_available()) # True device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # 实例化 Dataset # preprocessing normalize = transforms.Normalize(mean=[.5], std=[.5]) transform = transforms.Compose([transforms.ToTensor(), normalize]) # download and load the data train_dataset = torchvision.datasets.MNIST(root='./mnist/', train=True, transform=transform, download=True) test_dataset = torchvision.datasets.MNIST(root='./mnist/', train=False, transform=transform, download=False) # encapsulate them into dataloader form train_loader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, drop_last=True) test_loader = DataLoader(test_dataset, batch_size=args.batch_size, shuffle=False, drop_last=True) # ########### 保存检查点的地址(如果检查点不存在，则创建) ############ args.checkpoints_dir = os.path.join(args.train_log_dir, args.model_class + "_{}".format(args.train_tag), "checkpoints") if not os.path.exists(args.checkpoints_dir): os.makedirs(args.checkpoints_dir) ################################ # 实例化模型 # ################################ model = MNISTNet(input_dims=784, n_hiddens=256, n_class=10).to(device) # 实例化模型 ################################ # 损失函数 # ################################ ############################### # 创建优化器 Create optimizers # ############################### optimizer = torch.optim.Adam(model.parameters(), lr=args.lr) # ########### TensorBoard可视化 summary ############ event_dir = os.path.join(args.train_log_dir, args.model_class + "_{}".format(args.train_tag), "event_files") writer = SummaryWriter(event_dir) # 创建事件文件 # ########### 加载模型检查点 ############ start_epoch = 0 if args.model_name: print("加载模型：", args.checkpoints_dir + "/", args.model_name) checkpoint = torch.load(os.path.join(args.checkpoints_dir, args.model_name)) model.load_state_dict(checkpoint["model"]) optimizer.load_state_dict(checkpoint["optimizer"]) start_epoch = checkpoint['epoch'] for epoch in range(start_epoch, args.epochs): model.train() # 训练模型 correct = 0 for batch_idx, (data, target) in enumerate(train_loader): data = data.to(device) # torch.Size([64, 1, 28, 28]) target = target.to(device) # torch.Size([64]) optimizer.zero_grad() output = model(data) # torch.Size([64, 10]) loss = F.cross_entropy(output, target) loss.backward() optimizer.step() pred = output.data.max(1)[1] # 返回[值列表，索引列表]get the index of the max log-probability correct += pred.eq(target).sum() # ########### 可视化打印 ############ acc = 100. * correct / len(train_loader.dataset) print('Train Epoch: {} Loss: {:.6f} Accuracy: {:.4f}%'.format(epoch, loss, acc)) # ########### TensorBoard可视化 summary ############ # print('learning rate:', optimizer.state_dict()['param_groups'][0]['lr']) # writer.add_scalar(tag="lr", scalar_value=optimizer.state_dict()['param_groups'][0]['lr'], # global_step=epoch + 1) writer.add_scalar(tag="train/train_loss", scalar_value=loss, global_step=epoch + 1) writer.add_scalar(tag="train/train_acc", scalar_value=acc, global_step=epoch + 1) writer.flush() # 神经网络在测试验证集上的表现 model.eval() # 测试模型 test_loss = 0 correct = 0 # 测试的时候不需要梯度 with torch.no_grad(): for data, target in test_loader: data, target = data.to(device), target.to(device) output = model(data) # print(output.dtype, target.dtype) # torch.float32 torch.int64 test_loss += F.cross_entropy(output, target).data pred = output.data.max(1)[1] # get the index of the max log-probability correct += pred.eq(target).sum() # ########### 可视化打印 ############ test_loss = test_loss / len(test_loader) # average over number of mini-batch acc = 100. * correct / len(test_loader.dataset) print('\tTest set: Average loss: {:.4f}, Accuracy: {:.0f}%'.format(test_loss, acc)) ###################### # 更新tensorboard # ###################### writer.add_scalar(tag="val/test_loss", scalar_value=test_loss, global_step=epoch + 1) writer.add_scalar(tag="val/test_acc", scalar_value=acc, global_step=epoch + 1) writer.flush() # ########### 保存模型 ############ if (epoch + 1) % 10 == 0: print("保存模型") checkpoint = { "model": model.state_dict(), "optimizer": optimizer.state_dict(), "epoch": epoch + 1, } torch.save(checkpoint, '%s/%d.pth' % (args.checkpoints_dir, epoch + 1)) if __name__ == "__main__": main()
加载模型
深度学习的训练是需要很长时间的，我们不可能每次需要预测都花大量的时间去重新训练，因此我们想出一个方法，保存模型，也就是保存我们训练好的参数．
import numpy as np from keras.models import load_model from tensorflow.examples.tutorials.mnist import input_data mnist = input_data.read_data_sets("./mnist/", one_hot=True, reshape=False) # (?, 28,28,1) x_test = mnist.test.images # (10000, 28,28,1) y_test = mnist.test.labels # (10000, 10) print(y_test[1]) # [0. 0. 1. 0. 0. 0. 0. 0. 0. 0.] model = load_model('k_CNN.h5') # 读取模型 # 评估模型 evl = model.evaluate(x=x_test, y=y_test) evl_name = model.metrics_names for i in range(len(evl)): print(evl_name[i], ':\t', evl[i]) # loss : 0.19366768299341203 # acc : 0.9691 test = x_test[1].reshape(1, 28, 28, 1) y_predict = model.predict(test) # (1, 10) print(y_predict) # [[1.6e-06 6.0e-09 9.9e-01 5.8e-10 4.0e-07 2.5e-08 1.72e-06 1.2e-09 2.1e-07 8.5e-08]] y_true = 'actual:%d' % (np.argmax(y_test[1])) # actual:2 pre = 'actual:%d' % (np.argmax(y_predict)) # actual:2
参考文献
【博客园】MNIST数据集探究
【github】tensorflow-mnist-cnn
