这是我的完整代码

读取命令行参数

arg=argparse.ArgumentParser ()
arg.add_argument ("–num_epoch",default=5)# 可选命令行参数一般都要给默认值
arg.add_argument ("–batch_size",default=20)
arg.add_argument ("–learning_rate",default=0.001)
arg.add_argument ("–mode",default=“train”)
res=arg.parse_args ()

创建的模型类

class Model (tf.keras.Model):
# 创建模型中的所有层
def init(self):
super ().init()# 要调用 super () 来进行父类的初始化
self.flatten=tf.keras.layers.Flatten ()
self.dense1=tf.keras.layers.Dense (100,activation=tf.nn.relu)
self.dense2=tf.keras.layers.Dense (10,activation=None)
# 写清楚模型前向运算过程
@tf.function
def call (self,input):
x=self.flatten (input)
x=self.dense1 (x)
x=self.dense2 (x)
output=tf.nn.softmax (x)#one-hot 分类任务中一般在最后面都要加 softmax 层
return output

创建的数据类

class Data ():
# 该构造函数读取数据集，并对数据集进行归一化、扩维
def init(self):
(self.train_x,self.train_y),(self.test_x,self.test_y)=tf.keras.datasets.mnist.load_data ()
self.train_x=np.expand_dims (self.train_x.astype (np.float)/255.0,axis=-1)
self.test_x=np.expand_dims (self.test_x.astype (np.float)/255.0,axis=-1)
self.train_y=self.train_y.astype (np.int)
self.test_y=self.test_y.astype (np.int)
self.train_num=self.train_x.shape [0]
self.test_num=self.test_x.shape [0]
# 从训练集中每次采样 batch_size 的样本并返回
def get_data (self,batch_size):
vals=np.random.randint (0,self.train_num,batch_size)
return self.train_x [vals,:],self.train_y [vals]

训练函数

def train ():
# 读取命令行的超参数并且进行模型类和数据类的实例化
num_epochs=int (res.num_epoch)# 命令行传进来的默认为 str 类型，这里注意要把这三个值转化成对应的数值类型
batch_size=int (res.batch_size)
lr=float (res.learning_rate)
model=Model ()

# 创建模型保存器，这里加上 CheckpointManager () 的原因是可以设置一些自己的模型保存方式
# checkpoint=tf.train.Checkpoint (SModel=model)
# manager=tf.train.CheckpointManager (checkpoint,'./save',max_to_keep=3)
data=Data ()
optimizer=tf.keras.optimizers.Adam (learning_rate=lr)
num_size=data.train_num//batch_size*num_epochs
log='log'
# 创建写算子，方便后面将一些参数值写入日志文件
writer=tf.summary.create_file_writer (log)
# 开启日志追踪器，可以将计算图和一些操作的运算时间保存下来
tf.summary.trace_on (graph=True,profiler=True)
for i in range (num_size):
    x,y=data.get_data (batch_size)
    # 梯度计算器里面进行的是模型前向运算和损失函数计算的过程
    with tf.GradientTape () as tape:
        y_pre=model (x)
        loss=tf.keras.losses.sparse_categorical_crossentropy (y_true=y,y_pred=y_pre)
        loss=tf.reduce_sum (loss)
        # 将损失值写入日志文件
        with writer.as_default ():
            tf.summary.scalar ("loss",loss,step=i)
        if i%100==0:
            print ("训练批次为：{}  损失值为{}".format (i,loss))
    grad=tape.gradient (loss,model.variables)
    optimizer.apply_gradients (grads_and_vars=zip (grad,model.variables))
    # if i%100==0:
    #     path=manager.save (checkpoint_number=i)# 给保存的 ckpt 打上编号，并且返回保存路径
    #     print ("模型保存到了：{}".format (path))
# 保存 track 信息到日志文件
with writer.as_default ():
    tf.summary.trace_export ("trace",step=0,profiler_outdir=log)
tf.saved_model.save (model,'Trained_Models')
print ("模型保存完毕！")

训练函数

def test ():
# 往一个新的模型里面加载训练好的参数，这里是查找保存的最新的 ckpt 文件
# model=Model ()
# checkpoint=tf.train.Checkpoint (SModel=model)#SModel 名字要和保存时的名字相同
# checkpoint.restore (tf.train.latest_checkpoint (’./save’))
model=tf.saved_model.load (’./Trained_Models’)
print (“模型已经加载”)
data=Data ()
batch_size=int (res.batch_size)
num_size=data.test_num//batch_size
# 创建精度评估器
metri=tf.keras.metrics.SparseCategoricalAccuracy ()
for i in range (num_size):
x,y=data.test_x [i*batch_size:(i+1)*batch_size,:],data.test_y [i*batch_size:(i+1)*batch_size]
y_pre=model.call (x)
# 更新精度评估器
metri.update_state (y_true=y,y_pred=y_pre)
# 输出最终精度
print (“最终精度为：{}”.format (metri.result ()))
if name==‘main’:
if res.mode==‘train’:
train ()
if res.mode==‘test’:
test ()