Ggtcf

I'm benchmarking general TF operations, and so to establish a baseline I'm trying to figure out how quickly I can train a simple logistic regression with a single pass of the training data. My input is a TFRecord file containing 860,000 sparse rows, with 164,000 one-hot encoded features. Data processing at bottom.

A premade tf.estimator.Estimator, configured like so, can fit one pass of the data in 932 seconds :

feature_columns = [tf.feature_column.numeric_column(key='features',shape=164000)]

custom_config = tf.estimator.RunConfig(save_summary_steps=None,
 save_checkpoints_steps=None)

estimator = tf.estimator.LinearClassifier(
 feature_columns = feature_columns,
 model_dir = os.path.join(MODELDIR,f'PremadeLinearClassifier_currtime()'),
 config = custom_config
)

If I read the data into numpy arrays and create a dataset from from_tensor_slices() I can get that down to 467 seconds.

If I build my own training functions, I can perform a single pass of the data in 66.6 seconds, reading from disk.

class LogisticModel(object):
 def __init__(self):

 self.W = tf.Variable(tf.random_normal([164000,1],mean=0, stddev=0.1))
 self.B = tf.Variable(tf.random_normal([],mean=0.0,stddev=0.1))

 def __call__(self, x):

 return tf.sparse_tensor_dense_matmul(x,self.W) + self.B


def grad_fn(model, inputs, targets):

 with tf.GradientTape() as t:
 loss_val = loss_fn(model, inputs, targets)

 return t.gradient(loss_val, [model.W, model.B])


def loss_fn(model, inputs, targets):

 target_size = targets.shape.as_list()[0]
 return (tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(labels = tf.reshape(targets,[target_size,1]),
 logits = model(inputs))) 
 )

def perform_train(model, optim, dataset):

 step = 0
 for x, y in dataset:
 grads = grad_fn(lm, x, y)
 optimizer.apply_gradients(zip(grads, [lm.W, lm.B]),
 global_step = tf.train.get_or_create_global_step())

 if step % 20 == 0:
 print(f"Step step: loss_fn(lm,x,y)")
 step += 1

1) What could account for this huge speed difference? Is overhead in the estimators that significant?
2) Without reading into memory, can my custom function be improved further? An in-house C++-based library is still an order of magnitude faster.

Data generation:

def ex_to_tensors(ex, tensor_size):

 feature_spec = 'sparse': tf.SparseFeature(index_key='indices',
 value_key='values',
 dtype=tf.int64,
 size=tensor_size),
 'label': tf.FixedLenFeature([], tf.int64, default_value=0)
 

 parsed_dict = tf.parse_single_example(ex, feature_spec)

 return tf.cast(parsed_dict['sparse'],tf.float32), tf.cast(parsed_dict['label'],tf.float32)


def ex_input_fn(*filenames,batch_size=1000, feature_size=int(1e6)):

 def parseTensors(x):
 return ex_to_tensors(x,feature_size)

 dataset = (tf.data.TFRecordDataset(filenames)

 .map(parseTensors)
 .batch(batch_size)
 )

 return dataset