original_docstring: 
      Trains and export the model as SavedModel 2.0.
      Args:
        data_dir (str)           : Directory where to store datasets from TFDS.
                                   (With proper authentication, cloud bucket supported).
        buffer_size (int)        : Size of buffer to use while shuffling.
        batch_size (int)         : Size of each training batch.
        learning_rate (float)    : Learning rate to use for the optimizer.
        epoch (int)              : Number of Epochs to train for.
        dataset (tf.data.Dataset): Dataset object if data_dir is not provided.
        export_path (str)        : Path to export the trained model.
    
docstring: Trains and export the model as SavedModel 2.0.
docstring_tokens: Trains and export the model as SavedModel 2 . 0 .
code: def train_and_export(
          data_dir=None,
          buffer_size=1000,
          batch_size=32,
          learning_rate=1e-3,
          epoch=10,
          dataset=None,
          export_path="/tmp/tfhub_modules/mnist/digits/1"):
    """
      Trains and export the model as SavedModel 2.0.
      Args:
        data_dir (str)           : Directory where to store datasets from TFDS.
                                   (With proper authentication, cloud bucket supported).
        buffer_size (int)        : Size of buffer to use while shuffling.
        batch_size (int)         : Size of each training batch.
        learning_rate (float)    : Learning rate to use for the optimizer.
        epoch (int)              : Number of Epochs to train for.
        dataset (tf.data.Dataset): Dataset object if data_dir is not provided.
        export_path (str)        : Path to export the trained model.
    """
    model = MNIST()
    if not dataset:
      train = tfds.load(
          "mnist",
          split="train",
          data_dir=data_dir,
          batch_size=batch_size).shuffle(
          buffer_size,
          reshuffle_each_iteration=True)
    else:
      train = dataset
    optimizer_fn = tf.optimizers.Adam(learning_rate=learning_rate)
    loss_fn = tf.keras.losses.CategoricalCrossentropy(from_logits=True)
    metric = tf.keras.metrics.Mean()
    model.compile(optimizer_fn, loss=loss_fn)
    # Training Loop
    for epoch in range(epoch):
      for step, data in enumerate(train):
        train_step(
            model,
            loss_fn,
            optimizer_fn,
            metric,
            data['image'],
            data['label'])
        sys.stdout.write("\rEpoch: #{}\tStep: #{}\tLoss: {}".format(
            epoch, step, metric.result().numpy()))
    # Exporting Model as SavedModel 2.0
    tf.saved_model.save(model, export_path)
code_tokens: def train_and_export ( data_dir = None , buffer_size = 1000 , batch_size = 32 , learning_rate = 1e-3 , epoch = 10 , dataset = None , export_path = "/tmp/tfhub_modules/mnist/digits/1" ) : model = MNIST ( ) if not dataset : train = tfds . load ( "mnist" , split = "train" , data_dir = data_dir , batch_size = batch_size ) . shuffle ( buffer_size , reshuffle_each_iteration = True ) else : train = dataset optimizer_fn = tf . optimizers . Adam ( learning_rate = learning_rate ) loss_fn = tf . keras . losses . CategoricalCrossentropy ( from_logits = True ) metric = tf . keras . metrics . Mean ( ) model . compile ( optimizer_fn , loss = loss_fn ) for epoch in range ( epoch ) : for step , data in enumerate ( train ) : train_step ( model , loss_fn , optimizer_fn , metric , data [ 'image' ] , data [ 'label' ] ) sys . stdout . write ( "\rEpoch: #{}\tStep: #{}\tLoss: {}" . format ( epoch , step , metric . result ( ) . numpy ( ) ) ) tf . saved_model . save ( model , export_path )
short_docstring: Trains and export the model as SavedModel 2.0.
short_docstring_tokens: Trains and export the model as SavedModel 2 . 0 .
comment: """
      Trains and export the model as SavedModel 2.0.
      Args:
        data_dir (str)           : Directory where to store datasets from TFDS.
                                   (With proper authentication, cloud bucket supported).
        buffer_size (int)        : Size of buffer to use while shuffling.
        batch_size (int)         : Size of each training batch.
        learning_rate (float)    : Learning rate to use for the optimizer.
        epoch (int)              : Number of Epochs to train for.
        dataset (tf.data.Dataset): Dataset object if data_dir is not provided.
        export_path (str)        : Path to export the trained model.
    """$SEP$# Training Loop$SEP$# Exporting Model as SavedModel 2.0