AI开发平台MODELARTS-示例：从0到1制作自定义镜像并用于训练（Tensorflow+GPU）:Step3 准备训练脚本并上传至OBS

时间：2024-05-23 15:34:45

AI开发平台MODELARTS 模型训练

Step3 准备训练脚本并上传至OBS

准备本案例所需的训练脚本mnist.py，并上传至OBS桶的“obs://test-modelarts/tensorflow/code/”文件夹下。

mnist.py文件内容如下：

import argparse
import tensorflow as tf

parser = argparse.ArgumentParser(description='TensorFlow quick start')
parser.add_argument('--data_url', type=str, default="./Data", help='path where the dataset is saved')
args = parser.parse_args()

mnist = tf.keras.datasets.mnist

(x_train, y_train), (x_test, y_test) = mnist.load_data(args.data_url)
x_train, x_test = x_train / 255.0, x_test / 255.0

model = tf.keras.models.Sequential([
  tf.keras.layers.Flatten(input_shape=(28, 28)),
  tf.keras.layers.Dense(128, activation='relu'),
  tf.keras.layers.Dropout(0.2),
  tf.keras.layers.Dense(10)
])

loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)

model.compile(optimizer='adam',
              loss=loss_fn,
              metrics=['accuracy'])

model.fit(x_train, y_train, epochs=5)

上一篇：AI开发平台MODELARTS-示例：从0到1制作自定义镜像并用于训练（Tensorflow+GPU）:Step1 创建OBS桶和文件夹

下一篇：AI开发平台MODELARTS-示例：从0到1制作自定义镜像并用于训练（Tensorflow+GPU）:Step1 创建OBS桶和文件夹