0

我正在尝试使用混合精度来提高我的训练和推理性能，但我注意到它似乎并没有增加，实际上只是稍微降低了我的性能。我开始只是从Tensorflow网站上复制教程，即使这样似乎也没有任何变化。为了测试性能，我简单地计算了运行 fit 函数所需的时间。我正在使用 GeForce RTX 2080 Ti 和 tf 2.4.2。任何建议将不胜感激！

import tensorflow as tf

from tensorflow import keras
from tensorflow.keras import layers
from tensorflow.keras import mixed_precision
import time
mixed_precision.set_global_policy('mixed_float16')

# BUILD MODEL
inputs = keras.Input(shape=(784,), name='digits')
if tf.config.list_physical_devices('GPU'):
    print('The model will run with 4096 units on a GPU')
    num_units = 4096
else:
    # Use fewer units on CPUs so the model finishes in a reasonable amount of time
    print('The model will run with 64 units on a CPU')
    num_units = 64
dense1 = layers.Dense(num_units, activation='relu', name='dense_1')
x = dense1(inputs)
dense2 = layers.Dense(num_units, activation='relu', name='dense_2')
x = dense2(x)
x = layers.Dense(10, name='dense_logits')(x)
outputs = layers.Activation('softmax', dtype='float32', name='predictions')(x)
outputs = layers.Activation('linear', dtype='float32')(outputs)
model = keras.Model(inputs=inputs, outputs=outputs)
model.compile(loss='sparse_categorical_crossentropy',
              optimizer=keras.optimizers.RMSprop(),
              metrics=['accuracy'])

(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()
x_train = x_train.reshape(60000, 784).astype('float32') / 255
x_test = x_test.reshape(10000, 784).astype('float32') / 255
initial_weights = model.get_weights()

start = time.time()
history = model.fit(x_train, y_train,
                    batch_size=16384,
                    epochs=5,
                    validation_split=0.2)
end = time.time()
print(end-start)

4

0 回答 0