|
| 1 | +import tensorflow as tf |
| 2 | +from tensorflow.keras import datasets, layers, optimizers, Sequential, metrics |
| 3 | + |
| 4 | + |
| 5 | +# 设置GPU使用方式 |
| 6 | +# 获取GPU列表 |
| 7 | +gpus = tf.config.experimental.list_physical_devices('GPU') |
| 8 | +if gpus: |
| 9 | + try: |
| 10 | + # 设置GPU为增长式占用 |
| 11 | + for gpu in gpus: |
| 12 | + tf.config.experimental.set_memory_growth(gpu, True) |
| 13 | + except RuntimeError as e: |
| 14 | + # 打印异常 |
| 15 | + print(e) |
| 16 | + |
| 17 | +(xs, ys),_ = datasets.mnist.load_data() |
| 18 | +print('datasets:', xs.shape, ys.shape, xs.min(), xs.max()) |
| 19 | + |
| 20 | +batch_size = 32 |
| 21 | + |
| 22 | +xs = tf.convert_to_tensor(xs, dtype=tf.float32) / 255. |
| 23 | +db = tf.data.Dataset.from_tensor_slices((xs,ys)) |
| 24 | +db = db.batch(batch_size).repeat(30) |
| 25 | + |
| 26 | + |
| 27 | +model = Sequential([layers.Dense(256, activation='relu'), |
| 28 | + layers.Dense(128, activation='relu'), |
| 29 | + layers.Dense(10)]) |
| 30 | +model.build(input_shape=(4, 28*28)) |
| 31 | +model.summary() |
| 32 | + |
| 33 | +optimizer = optimizers.SGD(lr=0.01) |
| 34 | +acc_meter = metrics.Accuracy() |
| 35 | + |
| 36 | +for step, (x,y) in enumerate(db): |
| 37 | + |
| 38 | + with tf.GradientTape() as tape: |
| 39 | + # 打平操作,[b, 28, 28] => [b, 784] |
| 40 | + x = tf.reshape(x, (-1, 28*28)) |
| 41 | + # Step1. 得到模型输出output [b, 784] => [b, 10] |
| 42 | + out = model(x) |
| 43 | + # [b] => [b, 10] |
| 44 | + y_onehot = tf.one_hot(y, depth=10) |
| 45 | + # 计算差的平方和,[b, 10] |
| 46 | + loss = tf.square(out-y_onehot) |
| 47 | + # 计算每个样本的平均误差,[b] |
| 48 | + loss = tf.reduce_sum(loss) / x.shape[0] |
| 49 | + |
| 50 | + |
| 51 | + acc_meter.update_state(tf.argmax(out, axis=1), y) |
| 52 | + |
| 53 | + grads = tape.gradient(loss, model.trainable_variables) |
| 54 | + optimizer.apply_gradients(zip(grads, model.trainable_variables)) |
| 55 | + |
| 56 | + |
| 57 | + if step % 200==0: |
| 58 | + |
| 59 | + print(step, 'loss:', float(loss), 'acc:', acc_meter.result().numpy()) |
| 60 | + acc_meter.reset_states() |
0 commit comments