Skip to main content

Training utilities for TensorFlow.

Project description

Training utilities for TensorFlow.


pip install train

It is recommended to use a virtual environment.

Getting Started

from train import Model, GradientDescent
import tensorflow as tf

# Define the network architecture - layers, number of units, activations etc.
def network(inputs):
    hidden = tf.layers.Dense(units=64, activation=tf.nn.relu)(inputs)
    outputs = tf.layers.Dense(units=10)(hidden)
    return outputs

# Configure the learning process - loss, optimizer, evaluation metrics etc.
model = Model(network,

# Train the model using training data
model.train(x_train, y_train, epochs=30, batch_size=128)

# Evaluate the model performance on test or validation data
loss_and_metrics = model.evaluate(x_test, y_test)

# Use the model to make predictions for new data
predictions = model.predict(x)
# or call the model directly
predictions = model(x)

More configuration options are available:

model = Model(network,

You can also use custom functions for loss and metrics:

def custom_loss(labels, outputs):

def custom_metric(labels, outputs):

model = Model(network,
              metrics=['accuracy', custom_metric])

Model Function

To have more control, you may configure the model inside a function using Estimator class:

from train import Estimator, PREDICT
import tensorflow as tf

def model(features, labels, mode):
    # Define the network architecture
    hidden = tf.layers.Dense(units=64, activation=tf.nn.relu)(features)
    outputs = tf.layers.Dense(units=10)(hidden)
    predictions = tf.argmax(outputs, axis=1)
    # In prediction mode, simply return predictions without configuring learning process
    if mode == PREDICT:
        return predictions

    # Configure the learning process for training and evaluation modes
    loss = tf.losses.sparse_softmax_cross_entropy(labels, outputs)
    optimizer = tf.train.GradientDescentOptimizer(0.001)
    accuracy = tf.metrics.accuracy(labels, predictions)
    return dict(loss=loss,
                metrics={'accuracy': accuracy})

# Create the model using model function
model = Estimator(model)

# Train the model
model.train(x_train, y_train, epochs=30, batch_size=128)

mode parameter specifies whether the model is used for training, evaluation or prediction.

Training Mode

For layers like Dropout, you may use the training mode variable:

from train import training

x = tf.layers.Dropout(rate=0.4)(x, training=training())

Model and Estimator classes automatically manage the training mode variable. To change it manually, use:

# Enter training mode
training(True, session=session)

# Exit training mode
training(False, session=session)



Project details

Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

train-0.0.1.tar.gz (3.7 kB view hashes)

Uploaded Source

Supported by

AWS AWS Cloud computing and Security Sponsor Datadog Datadog Monitoring Fastly Fastly CDN Google Google Download Analytics Microsoft Microsoft PSF Sponsor Pingdom Pingdom Monitoring Sentry Sentry Error logging StatusPage StatusPage Status page