77 lines
2.2 KiB
Python
77 lines
2.2 KiB
Python
import tensorflow as tf
|
|
from tensorflow import keras
|
|
|
|
import numpy as np
|
|
import os
|
|
# Version Information
|
|
# tensorflow 2.2.0 , Cudnn7.6.5 and Cuda 10.1 , python 3.8
|
|
from keras import backend as K
|
|
K.clear_session()
|
|
|
|
gpus = tf.config.experimental.list_physical_devices('GPU')
|
|
if gpus:
|
|
try:
|
|
tf.config.experimental.set_virtual_device_configuration(gpus[0], [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=6024)])
|
|
except RuntimeError as e:
|
|
print(e)
|
|
os.exit(1)
|
|
|
|
print(tf.config.experimental.list_physical_devices())
|
|
|
|
print(tf.__version__)
|
|
|
|
print(tf.test.is_built_with_cuda())
|
|
|
|
(X_train, y_train), (X_test,y_test) = tf.keras.datasets.cifar10.load_data()
|
|
|
|
(X_train, y_train), (X_test,y_test) = tf.keras.datasets.cifar10.load_data()
|
|
|
|
print(X_train.shape,y_train.shape)
|
|
|
|
classes = ["airplane","automobile","bird","cat","deer","dog","frog","horse","ship","truck"]
|
|
|
|
print(classes[y_train[3][0]])
|
|
|
|
print("pre processing: scale images")
|
|
|
|
X_train_scaled = X_train / 255
|
|
X_test_scaled = X_test / 255
|
|
|
|
y_train_categorical = keras.utils.to_categorical(
|
|
y_train, num_classes=10, dtype='float32'
|
|
)
|
|
y_test_categorical = keras.utils.to_categorical(
|
|
y_test, num_classes=10, dtype='float32'
|
|
)
|
|
|
|
print("model build")
|
|
|
|
|
|
def get_model():
|
|
model = keras.Sequential([
|
|
keras.layers.Flatten(input_shape=(32,32,3)),
|
|
keras.layers.Dense(3000, activation='relu'),
|
|
keras.layers.Dense(1000, activation='relu'),
|
|
keras.layers.Dense(10, activation='sigmoid')
|
|
])
|
|
|
|
model.compile(optimizer='SGD',
|
|
loss='categorical_crossentropy',
|
|
metrics=['accuracy'])
|
|
return model
|
|
|
|
with tf.device('/GPU:0'):
|
|
model = keras.Sequential([
|
|
keras.layers.Flatten(input_shape=(32,32,3)),
|
|
keras.layers.Dense(3000, activation='relu'),
|
|
keras.layers.Dense(1000, activation='relu'),
|
|
keras.layers.Dense(10, activation='sigmoid')
|
|
])
|
|
|
|
model.compile(optimizer='SGD',
|
|
loss='categorical_crossentropy',
|
|
metrics=['accuracy'])
|
|
model.fit(X_train_scaled, y_train_categorical, epochs=10)
|
|
model.save('mymodel.keras')
|
|
|
|
print("done") |