resnet50_cifar10_rocm_benchmark

import tensorflow as tf
import numpy as np
gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
  try:
    for gpu in gpus:
      tf.config.experimental.set_memory_growth(gpu, True)
  except RuntimeError as e:
    print(e)

bsize = 500
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.cifar10.load_data()
x_train = (x_train/255).astype('float32')
x_test = (x_test/255).astype('float32')
n_classes = np.max(y_train)+1
train_dset = tf.data.Dataset.from_tensor_slices((x_train, y_train)).batch(bsize)
test_dset = tf.data.Dataset.from_tensor_slices((x_test, y_test)).batch(bsize)

mod = tf.keras.applications.resnet_v2.ResNet50V2(weights=None,
    input_shape = x_train.shape[1:],
    include_top = True,
    classes = n_classes,
    classifier_activation='softmax')
mod.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
#import pdb; pdb.set_trace()
mod.fit(train_dset, validation_data = test_dset, epochs = 8)