TensorFlow2.0仿pytorch快速搭建深度学习模型训练和文本分类任务实践_tensorflow2模型训练类似pytorch的写法

目录

一、TensorFlow2.0中的一些基本概念

1、张量Tensor

2、运算图graph

3、梯度正向和方向传播实现

4、数据管道DataSet

二、TensorFlow2.0仿pytorch搭建模型

三、TensorFlow2.0仿pytorch自定义模型训练

五、TensorFlow2.0仿pytorch实现文本分类任务实战

总结

参考文章

        作为一个NLP炼丹师，对当下比较流行的深度学习框架比较熟悉是很有必要的，至少要会用。pytorch比较好用，学习成本低，一般作为研究人员快速搭建模型验证算法的可行和正确性，对性能就不太苛求；而TensorFlow则是比较追求性能，但是1.0版本学习成本太高了。幸好TensorFlow2.0+版本问世，使得TensorFlow可以像pytorch那样快速构建模型，同时也能兼顾性能。这里主要对TensorFlow2.0如何像pytorch那样搭建一个模型和训练进行学习和总结，同时做一个分类任务的实战。

一、TensorFlow2.0中的一些基本概念

1、张量Tensor

这里的张量分为变量和常量，他们的区别是变量可以在计算图中用assign算子重新赋值。当然张量也是和numpy.array是一一对应的。

import tensorflow as tf
#常量张量
i = tf.constant(1) # tf.int32 类型常量
l = tf.constant(1,dtype = tf.int64) # tf.int64 类型常量
f = tf.constant(1.23) #tf.float32 类型常量
d = tf.constant(3.14,dtype = tf.double) # tf.double 类型常量
s = tf.constant("hello world") # tf.string类型常量
b = tf.constant(True) #tf.bool类型常量
#变量张量
v = tf.Variable([1.0,2.0],name = "v")
v.assign_add([1.0,1.0])

这里的张量操作和基本运算，以及维度的操作，都是和pytorch类似的；就是具体的API函数名不一样。例如：torch.cat()和tf.concat()、torch.mean()和tf.reduce_mean()，这里就不一一举例和说明了。

2、运算图graph

TensorFlow以前的版本中提供了静态计算图和动态计算图，它们都有自己的缺陷。那TensorFlow2.0就提供了一种Autograph计算图，它均衡了动态计算图和静态计算图的优势，削减他们的劣势。具体的实现就是采用装饰器，使用tf.function构建静态图的方式——Autograph。如下代码中，使用@tf.function把训练步骤中的一些计算放入静态图中计算，保证性能，同时构建的方式也很简单。

@tf.function
def train_step(model,input_ids,input_mask,labels,optimizer,train_loss,train_metric,loss_fun):
    with tf.GradientTape() as tape:
        predictions = model({'input_ids':input_ids,'attention_mask':input_mask})
        loss = loss_fun(labels,predictions)
    gradients = tape.gradient(loss,model.trainable_variables)
    optimizer.apply_gradients(zip(gradients,model.trainable_variables))
    train_loss.update_state(loss)
    train_metric.update_state(labels, predictions)

3、梯度正向和方向传播实现

torch中直接使用以下代码结构就可以实现这个功能

optimizer.zero_grad()
......
loss = F.cross_entropy(output,label)
loss.backward()
optimizer.step()

类似的TensorFlow2.0中也提供了类似的这样一种代码结构，最重要的就是tf.GradientTape()梯度磁带来记录正向计算，然后使用反向传播磁带自动获得梯度值；结合优化器就可以优化网络了。

with tf.GradientTape() as tape:
   predictions = model(inputs)#得到输出
   loss = loss_fun(labels,predictions)#得到loss
gradients = tape.gradient(loss,model.trainable_variables)#得到梯度
optimizer.apply_gradients(zip(gradients,model.trainable_variables))#梯度优化
train_loss.update_state(loss)#loss更新
train_metric.update_state(labels, predictions)#评价指标更新

4、数据管道DataSet

torch中的数据管道，只需要把数据读取到管道中，然后在调用torch.utils.data.DataLoader就能把数据安装batch一分一分的传入模型中进行训练和迭代。TensorFlow2.0中也提供了各式各样的数据管道，可以从 Numpy array, Pandas DataFrame, Python generator, csv，tfrecords等文件以及路径构建数据管道。tfrecords这种比较麻烦，但是优点明显——压缩后文件较小，便于网络传播，加载速度较快，目前我还不怎么会用，后面有机会学习。这里就给最简单的做个实例，tf.data.Dataset.from_generator()是比较简单和常用的一种。注意的是这里一定要是参数是迭代器，同时也要指定输出的数据类型。

data = tf.data.Dataset.from_generator(generator,output_types=(tf.float32,tf.int32))

这里还没有对数据进行乱序和batch处理，要处理就得添加对应的代码：

data = data.shuffle(buffer_size=1000).batch(args.batch_size)

这里还有一个比较好的函数 prefetch (),它可以让数据准备和参数迭代两个过程相互并行。

比较规范和标准的用法就是：

train_data = tf.data.Dataset.from_generator(train_generator,(tf.int32,tf.int32,tf.int32)).shuffle(buffer_size=1000).batch(args.batch_size).prefetch(buffer_size = tf.data.experimental.AUTOTUNE)

简单的说明一下，这里buffer_size的作用，它就是限定内存中用来打乱顺序数据的条数以及数据准备和参数迭代并行的条数。

二、TensorFlow2.0仿pytorch搭建模型

          TensorFlow2.0中提供了多种多样的模型构建方式，1使用Sequential按层顺序构建模型,2使用函数式API构建任意结构模型，3继承Model基类构建自定义模型。其中我个人最喜欢的就是第三种继承Model基类构建自定义模型，因为我比较熟悉pytorch那一套构建模型的流程。当然这三种方式构建模型的方法各有优点，同时在最后模型的保存过程中也是采用了不同的方式，这个可以总结一个博客出来，后面内容部分做一个简单的介绍。

           直接看第三种继承Model基类构建自定义模型是什么样的。继承Model基类的自定义模型类中需要有2个函数，一个是init函数，一个是call函数。当然有的也会有build函数，其实这个build函数的内容是可以直接和init合并的。上代码，先上torch版本的，然后在上TensorFlow版本的。

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import BertModel
class TextBert(nn.Module):
    def __init__(self,args):
        super(TextBert,self).__init__()
        self.bert = BertModel.from_pretrained(args.model_path)
        self.cl1 = nn.Linear(768,768)
        self.cl2 = nn.Linear(768,384)
        self.cl3 = nn.Linear(384, 8)
 
    def forward(self,input_ids,input_mask):
        embedding = self.bert(input_ids,input_mask)[0]
        mean_embedding = torch.mean(embedding,dim=1)
        x = self.dropout(mean_embedding)
        x = F.relu(self.cl1(x))
        x = F.relu(self.cl2(x))
        logit = self.cl3(x)
        return logit

TensorFlow版本的如下：

import tensorflow as tf
from tensorflow.keras import *
from transformers import TFBertModel
 
class TextBert(models.Model):
    def __init__(self,args,):
        super(TextBert,self).__init__()
        self.args = args
        #由于没有下载到model.h5模型权重，这里得使用pytorch_model.bin格式文件加载from_pt = true
        self.bert = TFBertModel.from_pretrained(self.args.model_path, from_pt=True)
        self.cl1 = layers.Dense(768, activation='relu')
        self.cl2 = layers.Dense(384, activation='relu')
        self.cl3 = layers.Dense(8, activation='softmax')
 
 
    def call(self, inputs):
        embedding = self.bert(inputs)[0]
        mean_embedding = tf.reduce_mean(embedding,1)
        x = self.cl1(mean_embedding)
        x = tf.nn.relu(x)
        x = self.cl2(x)
        x = tf.nn.relu(x)
        logit = self.cl3(x)
        return logit

可以看到torch和TensorFlow2.0框架构建模型流程类似，torch中forward()和tf中call()有相同的作用；然后torch继承的是nn.Module而TensorFlow继承的是models.Model；其他的子模块相同的功能大都可以简单的理解为不同的函数名具有相似的作用，所以这个模型构建起来就很快速高效了。

最后说一下build这个函数，在模型中没有这个函数的话，就不能在构建模型后主动的显式调用build()函数，然后调用model.summry()来查看模型的具体结构和参数了。当然也是可以经历过一个输入后，然后调用model.summry()就可以查看模型的结构和参数量了。

三、TensorFlow2.0仿pytorch自定义模型训练

模型构建好了以后，就可以输入数据集，采用合适的优化器和损失函数来进行训练了。torch中，流程是这样的：定义优化器和损失函数，学习率等等；优化器梯度清零，模型输出和label得到具体的loss，loss反向传播，优化器更新梯度；模型性能评估，训练集验证集准确率等等。同样的TensorFlow中也可以这样来写，只不过代码细节可能不一样。这里给出一个简单的训练过程：

def train_model(model,train_data,train_len,dev_data,args):
    #定义优化器，loss和评估指标
    optimizer = optimizers.Adam(learning_rate=args.lr)
    train_loss = metrics.Mean(name='train_loss')
    train_metric = metrics.SparseCategoricalAccuracy(name='train_accuracy')
 
    valid_loss = metrics.Mean(name='valid_loss')
    valid_metric = metrics.SparseCategoricalAccuracy(name='valid_accuracy')
 
    #定义损失函数
    loss_fun = losses.SparseCategoricalCrossentropy()
    step = 0
    best_valid_acc = 0
    for epoch in tf.range(args.epochs):
        for input_ids,input_mask,labels in train_data:
            #训练模型
            train_step(model,input_ids,input_mask,labels,optimizer,train_loss,train_metric,loss_fun)
            step += 1
            # print('step',step)
            if step%100 ==0 and step%((int(train_len/args.batch_size/2/100))*100)!=0:
                logs = 'Epoch={},step={},Loss:{},Accuracy:{},Valid Loss:{},Valid Accuracy:{},best_valid_acc:{}'
                tf.print(tf.strings.format(logs, (
                    epoch, step, train_loss.result(), train_metric.result(), valid_loss.result(), valid_metric.result(),
                    best_valid_acc)))
 
            if step%((int(train_len/args.batch_size/2/100))*100)==0:
                for input_ids, input_mask, labels in dev_data:
                    #验证模型
                    valid_step(model, input_ids, input_mask, labels, optimizer, valid_loss, valid_metric,loss_fun)
                if valid_metric.result()>=best_valid_acc:
                    best_valid_acc = valid_metric.result()
                    save_path = args.model_save_path
                    # model.save(save_path,save_format='h5')
                    # model.save(save_path, save_format='tf')
                    #保存模型
                    model.save_weights(save_path,save_format='tf')
                logs = 'Epoch={},step={},Loss:{},Accuracy:{},Valid Loss:{},Valid Accuracy:{},best_valid_acc:{}'
                printbar()
                tf.print(tf.strings.format(logs, (
                epoch,step,train_loss.result(), train_metric.result(), valid_loss.result(), valid_metric.result(),best_valid_acc)))
                tf.print("")
 
        #loss和metric更新
        train_loss.reset_states()
        train_metric.reset_states()
        valid_loss.reset_states()
        valid_metric.reset_states()

具体的训练细节：

@tf.function
def train_step(model,input_ids,input_mask,labels,optimizer,train_loss,train_metric,loss_fun):
    with tf.GradientTape() as tape:
        #模型输出
        predictions = model({'input_ids':input_ids,'attention_mask':input_mask})
        #label和模型输出得到loss
        loss = loss_fun(labels,predictions)
    #得到参数的梯度
    gradients = tape.gradient(loss,model.trainable_variables)
    #优化器更新梯度
    optimizer.apply_gradients(zip(gradients,model.trainable_variables))
    train_loss.update_state(loss)
    train_metric.update_state(labels, predictions)

还是和torch有一些不同。

TensorFlow2.0中自定义模型的训练过程差不多就是上面这么多内容了，可以扩展的地方就是优化器的选择，loss函数的选择以及学习率的衰减策略、权重冻结；不同模型层采用不同的学习率的实现等等。

四、模型保存和加载

TensorFlow中模型的保存有多种方式，这就与模型是通过那种方式构建的有密切关系，要是不对应的话就会报错。

采用sequential和函数式API的构建方式，就可以保存为.h5格式的文件，也可以保存为savedModel模型。实例如下：

"""
sequential模型可以保存为完整的.h5模型和savedModel——(这是一个文件夹）
"""
 
import tensorflow as tf
from tensorflow.keras import models,layers
import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '1'
def auto_distribute_gpu_memory():
    gpus = tf.config.experimental.list_physical_devices("GPU")
    if gpus:
        try:
            for gpu in gpus:
                tf.config.experimental.set_memory_growth(gpu,True)
        except RuntimeError as e:
            print(e)
 
def create_model():
    model = models.Sequential()
    model.add(layers.Embedding(200,768,input_length=512))
    model.add(layers.Conv1D(filters=64,kernel_size=5,activation='relu'))
    model.add(layers.MaxPool1D(2))
    model.add(layers.Flatten())
    model.add(layers.Dense(2, activation="softmax"))
    return model
 
if __name__ == '__main__':
    auto_distribute_gpu_memory()
    seqential_model = create_model()
    seqential_model.compile(optimizer='Nadam',
            loss='binary_crossentropy',
            metrics=['accuracy',"AUC"])
    seqential_model.summary()
    seqential_model.save('seqential_model.h5')
    seqential_model.save('SavedModel',save_format='tf')
 
    #保存完整的模型结构和参数
    h5_model = models.load_model('seqential_model.h5')
    h5_model.summary()
 
    savedModel = models.load_model('SavedModel')
    savedModel.summary()

注意采用的API就是如下代码，要配套使用

#模型保存
seqential_model.save('seqential_model.h5')
seqential_model.save('SavedModel',save_format='tf')
#模型加载
models.load_model('seqential_model.h5')
models.load_model('SavedModel')

当采用自定义继承models.Model来实现模型的构建的时候，就不能存取为.h5格式的文件，它可以被保存为SavedModel类型和checkpoint类型。下面只展示SavedModel类型的保存。这里还有个点比较重要，就是模型初始化以后，如果没有被"使用"保存的过程中仍然会报错，同时如果没有build，调用summary()也会报错。

"""
sequential模型可以保存为完整的.h5模型和savedModel——(这是一个文件夹）
"""
 
import tensorflow as tf
from tensorflow.keras import models,layers
import os
import numpy as np
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '1'
def auto_distribute_gpu_memory():
    gpus = tf.config.experimental.list_physical_devices("GPU")
    if gpus:
        try:
            for gpu in gpus:
                tf.config.experimental.set_memory_growth(gpu,True)
        except RuntimeError as e:
            print(e)
 
class Mymodel(models.Model):
    def __init__(self):
        super(Mymodel,self).__init__()
        self.emd = layers.Embedding(200,768)
        self.conv = layers.Conv1D(filters=64,kernel_size=5,activation='relu')
        self.pool = layers.MaxPool1D(2)
        self.flat = layers.Flatten()
        self.cl = layers.Dense(2, activation="softmax")
 
 
    def call(self, inputs, training=None, mask=None):
        x = self.emd(inputs)
        x = self.conv(x)
        x = self.pool(x)
        x = self.flat(x)
        output = self.cl(x)
        return output
 
if __name__ == '__main__':
    auto_distribute_gpu_memory()
    sub_class_model = Mymodel()
    sub_class_model.compile(optimizer='Nadam',
            loss='binary_crossentropy',
            metrics=['accuracy',"AUC"])
    #没有build的话，summary()的调用也会出错的
    sub_class_model.build(input_shape=(None,512))
    sub_class_model.summary()
 
    #TensorSpec创建一个“无实际数据的张量”，指定它的形状，作为模型的输入
    shape = tf.TensorSpec(shape=(100, 512), dtype=tf.dtypes.int32, name=None)
    #没有_set_inputs这一步模型保存会报错
    sub_class_model._set_inputs(shape)
    sub_class_model.save('sub_class_model',save_format='tf')
 
    #保存完整的模型结构和参数
    sub_class_h5_model = models.load_model('sub_class_model')
    sub_class_h5_model.summary()

没有正确保存模型的结果如下图：

五、TensorFlow2.0仿pytorch实现文本分类任务实战

这里主要是把我以前实现过的基于Bert和pytorch长文本的分类任务，用TensorFlow2.0来实现一遍，原项目情况：基于机器学习算法和pytorch实现的深度学习模型的中文长文本多分类任务实战

首先看项目目录：

这个和torch那一套流程类似。

第一个改变的就是ReadDataSet.py，由于把文本转化为tensor后分批喂入模型中，采用的是tf的数据管道，这个和torch的DataLoader是不同的，因此ReadDataSet.py中要修改相应代码。这里采用的数据管道就是生成tf.data.Dataset.from_generator()这种格式。所以在ReadDataSet.py中要返回一个生成器。关键代码如下：

def __call__(self, *args, **kwargs):
     for (input_ids,input_mask,label) in self.tf_tensors_list:
          yield input_ids,input_mask,label

全部的ReadDataSet.py代码如下，主要是实现文本数据通过BertTokenizer做一个隐射，包装成tf张量，同时生成一个生成器。

from transformers import BertTokenizer
from tqdm import tqdm
import os
import logging
import tensorflow as tf
 
logger = logging.getLogger(__name__)
 
class ReadDataSet():
    def __init__(self,data_file_name,args,repeat=1):
        self.max_sentence_length = args.max_sentence_length
        self.repeat = repeat
        self.tokenizer = BertTokenizer.from_pretrained(args.model_path)
        self.file_path = args.data_file_path
        self.file_name = data_file_name
        self.tf_tensors_list = self.readfiles_tokening()
        if 'train' in self.file_name:
            self.len_train = len(self.tf_tensors_list)
 
 
    def readfiles_tokening(self):
        file_path = os.path.join(self.file_path, self.file_name)
        tf_tensors_list = []
        with open(file_path, 'r') as f:
            lines = f.readlines()
        for line in tqdm(lines, desc='read data tokening'):
            line = line.strip().split('\t')
            input_ids,input_mask,label = self.do_process_data((line[0], line[1]))
            tf_tensors_list.append((input_ids,input_mask,label))
        return tf_tensors_list
 
 
    def __call__(self, *args, **kwargs):
        for (input_ids,input_mask,label) in self.tf_tensors_list:
            yield input_ids,input_mask,label
 
    def do_process_data(self, params):
        sentence = params[0]
        label = params[1]
        input_ids, input_mask = self.convert_into_indextokens_and_segment_id(sentence)
        input_ids = tf.constant(input_ids)
        input_mask = tf.constant(input_mask)
        label = tf.constant(label)
        return input_ids,input_mask,label
 
    def convert_into_indextokens_and_segment_id(self, text):
        tokeniz_text = self.tokenizer.tokenize(text[0:self.max_sentence_length])
        input_ids = self.tokenizer.convert_tokens_to_ids(tokeniz_text)
        input_mask = [1] * len(input_ids)
 
        pad_indextokens = [0] * (self.max_sentence_length - len(input_ids))
        input_ids.extend(pad_indextokens)
        input_mask_pad = [0] * (self.max_sentence_length - len(input_mask))
        input_mask.extend(input_mask_pad)
        return input_ids, input_mask
 

模型代码就比较简单，和torch构建的模型差不太多

import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '1'
import tensorflow as tf
from tensorflow.keras import *
from transformers import TFBertModel
 
"""
这里的模型设置的是bert模型在训练的过程中不会改变权重，这个可以和bert权重参与训练做对比
"""
 
class TextBert(models.Model):
    def __init__(self,args,):
        super(TextBert,self).__init__()
        self.args = args
        #由于没有下载到model.h5模型权重，这里得使用pytorch_model.bin格式文件加载from_pt = true
        self.bert = TFBertModel.from_pretrained(self.args.model_path, from_pt=True)
        self.cl1 = layers.Dense(768, activation='relu')
        self.cl2 = layers.Dense(384, activation='relu')
        self.cl3 = layers.Dense(8, activation='softmax')
 
 
    def call(self, inputs):
        embedding = self.bert(inputs)[0]
        mean_embedding = tf.reduce_mean(embedding,1)
        x = self.cl1(mean_embedding)
        x = tf.nn.relu(x)
        x = self.cl2(x)
        x = tf.nn.relu(x)
        logit = self.cl3(x)
        return logit
 
 

这里有一个值得注意的是，huggineface并没有提供中文的roberta模型的权重模型.h5格式的文件，所以采用pytorch_model.bin格式文件加载from_pt = true。

训练代码也没有什么特别注意的地方，上面多已经说了怎么自定义模型训练，直接上代码;

import tensorflow as tf
from tensorflow.keras import optimizers,metrics,losses
from tqdm import tqdm
from model.TextBert import TextBert
import argparse
from DataReader.ReadDataSet import ReadDataSet
 
import os
 
os.environ["CUDA_VISIBLE_DEVICES"] = "0"       # 设置使用哪一块GPU（默认是从0开始）
 
# 下面就是实现按需分配的代码！
gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
  try:
    # Currently, memory growth needs to be the same across GPUs
    for gpu in gpus:
      tf.config.experimental.set_memory_growth(gpu, True)
    logical_gpus = tf.config.experimental.list_logical_devices('GPU')
    print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
  except RuntimeError as e:
    # Memory growth must be set before GPUs have been initialized
    print(e)
 
def printbar():
    today = tf.timestamp()%(24*60*60)
 
    hour = tf.cast(today // 3600 + 8, tf.int32) % tf.constant(24)
    minite = tf.cast((today % 3600) // 60, tf.int32)
    second = tf.cast(tf.floor(today % 60), tf.int32)
    def timeformat(m):
        if tf.strings.length(tf.strings.format("{}", m))==1:
            return (tf.strings.format("0{}", m))
        else:
            return (tf.strings.format("{}", m))
 
    timestrins = tf.strings.join([timeformat(hour), timeformat(minite),
                                  timeformat(second)], separator=":")
    tf.print('====' * 20 + timestrins)
@tf.function
def train_step(model,input_ids,input_mask,labels,optimizer,train_loss,train_metric,loss_fun):
    with tf.GradientTape() as tape:
        #模型输出
        predictions = model({'input_ids':input_ids,'attention_mask':input_mask})
        #label和模型输出得到loss
        loss = loss_fun(labels,predictions)
    #得到参数的梯度
    gradients = tape.gradient(loss,model.trainable_variables)
    #优化器更新梯度
    optimizer.apply_gradients(zip(gradients,model.trainable_variables))
    train_loss.update_state(loss)
    train_metric.update_state(labels, predictions)
 
def valid_step(model, input_ids, input_mask, labels, optimizer, valid_loss, valid_metric,loss_fun):
    predictions = model({'input_ids':input_ids,'attention_mask':input_mask})
    batch_loss = loss_fun(labels, predictions)
    valid_loss.update_state(batch_loss)
    valid_metric.update_state(labels, predictions)
 
def train_model(model,train_data,train_len,dev_data,args):
    #定义优化器，loss和评估指标
    optimizer = optimizers.Adam(learning_rate=args.lr)
    train_loss = metrics.Mean(name='train_loss')
    train_metric = metrics.SparseCategoricalAccuracy(name='train_accuracy')
 
    valid_loss = metrics.Mean(name='valid_loss')
    valid_metric = metrics.SparseCategoricalAccuracy(name='valid_accuracy')
 
    #定义损失函数
    loss_fun = losses.SparseCategoricalCrossentropy()
    step = 0
    best_valid_acc = 0
    for epoch in tf.range(args.epochs):
        for input_ids,input_mask,labels in train_data:
            #训练模型
            train_step(model,input_ids,input_mask,labels,optimizer,train_loss,train_metric,loss_fun)
            step += 1
            # print('step',step)
            if step%100 ==0 and step%((int(train_len/args.batch_size/2/100))*100)!=0:
                logs = 'Epoch={},step={},Loss:{},Accuracy:{},Valid Loss:{},Valid Accuracy:{},best_valid_acc:{}'
                tf.print(tf.strings.format(logs, (
                    epoch, step, train_loss.result(), train_metric.result(), valid_loss.result(), valid_metric.result(),
                    best_valid_acc)))
 
            if step%((int(train_len/args.batch_size/2/100))*100)==0:
                for input_ids, input_mask, labels in dev_data:
                    #验证模型
                    valid_step(model, input_ids, input_mask, labels, optimizer, valid_loss, valid_metric,loss_fun)
                if valid_metric.result()>=best_valid_acc:
                    best_valid_acc = valid_metric.result()
                    save_path = args.model_save_path
                    # model.save(save_path,save_format='h5')
                    # model.save(save_path, save_format='tf')
                    #保存模型
                    model.save_weights(save_path,save_format='tf')
                logs = 'Epoch={},step={},Loss:{},Accuracy:{},Valid Loss:{},Valid Accuracy:{},best_valid_acc:{}'
                printbar()
                tf.print(tf.strings.format(logs, (
                epoch,step,train_loss.result(), train_metric.result(), valid_loss.result(), valid_metric.result(),best_valid_acc)))
                tf.print("")
 
        #loss和metric更新
        train_loss.reset_states()
        train_metric.reset_states()
        valid_loss.reset_states()
        valid_metric.reset_states()
 
if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='init params configuration')
    parser.add_argument('--batch_size',type=int,default=100)
    parser.add_argument('--model_path',type=str,default='./pretrain_model/')
    parser.add_argument('--requires_grad', type= bool,default=True)
    parser.add_argument('--data_file_path',type=str,default='data_set/patent')
    parser.add_argument('--max_sentence_length',type=int,default=400)
    parser.add_argument('--lr', type=float, default=0.00001)
    parser.add_argument('--epochs', type=int, default=3)
    parser.add_argument('--model_save_path', type=str, default='savedmodel/')
    args = parser.parse_args()
    print(args)
    train_generator = ReadDataSet('train.tsv',args)
    train_len = train_generator.len_train
    print('*'*80,train_len)
    train_data = tf.data.Dataset.from_generator(train_generator,(tf.int32,tf.int32,tf.int32)).shuffle(buffer_size=1000).batch(args.batch_size).prefetch(buffer_size = tf.data.experimental.AUTOTUNE)
    dev_generator = ReadDataSet('dev.tsv', args)
    dev_data = tf.data.Dataset.from_generator(dev_generator, (tf.int32, tf.int32, tf.int32)).shuffle(buffer_size=1000).batch(args.batch_size).prefetch(buffer_size = tf.data.experimental.AUTOTUNE)
    model = TextBert(args)
    train_model(model,train_data,train_len,dev_data,args)
 

训练过程和结果如下：

由于采用的是roberta模型得到验证集的准确率是83.77%，比以前bert+pytorch的83.59%高了0.18%。当然这里采用的max_sen_length也是510，比以前工作的400要长一点。

总结

本博客首先对TensorFlow2.0中的一些基本概念做了展示说明，然后分别对模型的构建和自定义训练以及模型的保存和加载进行描述，最后用一个分类任务，对TensorFlow2.0的使用进行实践。当然TensorFlow这个框架是有很多地方值得学习的，由于能力有限和学习实践不长，这里只能进行一个常用模块和模型建立训练流程做一个学习，更加细节的地方，留着以后慢慢研究学习。

参考文章

详解tensorflow2.0的模型保存方法（一）

eat_tensorflow2_in_30_days