How to build an age and gender multi-task predictor with deep learning in TensorFlow


In my last tutorial, you learned about how to combine a convolutional neural network and Long short-term memory (LTSM) to create captions given an image. In this tutorial, you’ll learn how to build and train a multi-task machine learning model to predict the age and gender of a subject in an image.


  • Introduction to age and gender model
  • Building a Multi-task Tensorflow Estimator
  • Training


  • basic understanding of convolutional neural networks (CNN)
  • basic understanding of TensorFlow
  • GPU (optional)

Introduction to Age and Gender Model

In 2015, researchers from Computer Vision Lab, D-ITET, published a paper DEX and made public their IMDB-WIKI consisting of 500K+ face images with age and gender labels.

IMDB-WIKI Dataset source:

DEX outlines an neural network architecture involving a pretrained imagenet vgg16 model that estimates the apparent age in face images. DEX placed first in ChaLearn LAP 2015 — a competition that deals with recognizing people in an image — outperforming human reference.

Age as a classification problem

A conventional way of tackling an age estimation problem with an image as input would be using a regression-based model with mean-squared error as the loss function. DEX models this problem as a classification task, using a softmax classifier with each age represented as a unique class ranging from 1 to 101 and cross-entropy as the loss function.

Multi-task learning

Multi-task learning is a technique of training on multiple tasks through a shared architecture. Layers at the beginning of the network will learn a joint generalized representation, preventing overfitting to a specific task that may contain noise.

By training with a multi-task network, the network can be trained in parallel on both tasks. This reduces the infrastructure complexity to only one training pipeline. Additionally, the computation required for training is reduced as both tasks are trained simultaneously.

Multi-task CNN source:

Building a multi-task network in TensorFlow

Below you’ll use TensorFlow’s estimator abstraction to create the model. The model will be trained from raw image input to predict the age and gender of the face image.

Project Structure

    ├── Dockerfile
    ├── age_gender_estimation_tutorial
    │   ├──
    │   ├──
    │   └──
    ├── bin
    │   ├──
    │   ├──
    │   ├──
    │   └──
    ├── requirements.txt


For the environment, you’ll use Docker to install dependencies. A GPU version is also provided for convenience.

FROM tensorflow/tensorflow:1.12.0-py3
RUN apt-get update \
    && apt-get install -y libsm6 libxrender-dev libxext6
ADD $PWD/requirements.txt /requirements.txt
RUN pip3 install -r /requirements.txt
CMD ["/bin/bash"]

Dockerfile (cpu version)

FROM tensorflow/tensorflow:1.12.0-gpu-py3
RUN apt-get update \
    && apt-get install -y libsm6 libxrender-dev libxext6
ADD $PWD/requirements.txt /requirements.txt
RUN pip3 install -r /requirements.txt
CMD ["/bin/bash"]

Dockerfile.gpu (GPU version)



docker build -t colemurray/age-gender-estimation-tutorial -f Dockerfile .


To train this model, you’ll use the IMDB-WIKI dataset, consisting of 500K+ images. For simplicity, you’ll download the pre-cropped imdb images (7GB). Run the script below to download the data.

#!/usr/bin/env bash
if [[ ! -d "data" ]]
mkdir "data"
curl -O
tar -xzvf imdb_crop -C data

    chmod +x bin/



You’ll now process the dataset to clean out low-quality images and crop the input to a fixed image size. Additionally, you’ll format the data as a CSV to simplify reading into TensorFlow.

import argparse as argparse
import csv
import os
import random
from datetime import datetime
import cv2
import numpy as np
from import loadmat
from tqdm import tqdm
headers = ['filename', 'age', 'gender']
def calc_age(taken, dob):
    birth = datetime.fromordinal(max(int(dob) - 366, 1))
    # assume the photo was taken in the middle of the year
    if birth.month < 7:
        return taken - birth.year
        return taken - birth.year - 1
def load_db(mat_path):
    db = loadmat(mat_path)['imdb'][0, 0]
    num_records = len(db["face_score"][0])
    return db, num_records
def get_meta(db):
    full_path = db["full_path"][0]
    dob = db["dob"][0]  # Matlab serial date number
    gender = db["gender"][0]
    photo_taken = db["photo_taken"][0]  # year
    face_score = db["face_score"][0]
    second_face_score = db["second_face_score"][0]
    age = [calc_age(photo_taken[i], dob[i]) for i in range(len(dob))]
    return full_path, dob, gender, photo_taken, face_score, second_face_score, age
def main(input_db, photo_dir, output_dir, min_score=1.0, img_size=165, split_ratio=0.8):
    Takes imdb dataset db and performs processing such as cropping and quality checks, writing output to a csv.
    :param split_ratio:
    :param input_db: Path to imdb db
    :param photo_dir: Path to photo's directory
    :param output_dir: Directory to write output to
    :param min_score: minimum score to filter face quality, range [0, 1.0]
    :param img_size: size to crop images to
    crop_dir = os.path.join(output_dir, 'crop')
    if not os.path.exists(output_dir):
    if not os.path.exists(crop_dir):
    db, num_records = load_db(input_db)
    indices = list(range(num_records))
    train_indices = indices[:int(len(indices) * split_ratio)]
    test_indices = indices[int(len(indices) * split_ratio):]
    train_csv = open(os.path.join(output_dir, 'train.csv'), 'w')
    train_writer = csv.writer(train_csv, delimiter=',', )
    val_csv = open(os.path.join(output_dir, 'val.csv'), 'w')
    val_writer = csv.writer(val_csv, delimiter=',')
    clean_and_resize(db, photo_dir, train_indices, min_score, img_size, train_writer, crop_dir)
    clean_and_resize(db, photo_dir, test_indices, min_score, img_size, val_writer, crop_dir)
def clean_and_resize(db, photo_dir, indices, min_score, img_size, writer, crop_dir):
    Cleans records and writes output to :param writer
    :param db:
    :param photo_dir:
    :param indices:
    :param min_score:
    :param img_size:
    :param crop_dir:
    :param writer:
    full_path, dob, gender, photo_taken, face_score, second_face_score, age = get_meta(db)
    for i in tqdm(indices):
        filename = str(full_path[i][0])
        if not os.path.exists(os.path.join(crop_dir, os.path.dirname(filename))):
            os.makedirs(os.path.join(crop_dir, os.path.dirname(filename)))
        img_path = os.path.join(photo_dir, filename)
        if float(face_score[i]) < min_score:
        if (~np.isnan(second_face_score[i])) and second_face_score[i] > 0.0:
        if ~(0 <= age[i] <= 100):
        if np.isnan(gender[i]):
        img_gender = int(gender[i])
        img_age = int(age[i])
        img = cv2.imread(img_path)
        crop = cv2.resize(img, (img_size, img_size))
        crop_filepath = os.path.join(crop_dir, filename)
        cv2.imwrite(crop_filepath, crop)
        writer.writerow([filename, img_age, img_gender])
if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--db-path', required=True)
    parser.add_argument('--photo-dir', required=True)
    parser.add_argument('--output-dir', required=True)
    parser.add_argument('--min-score', required=False, type=float, default=1.0)
    parser.add_argument('--img-size', type=int, required=False, default=224)
    parser.add_argument('--split-ratio', type=float, required=False, default=0.8)
    args = parser.parse_args()
    main(input_db=args.db_path, photo_dir=args.photo_dir, output_dir=args.output_dir,
         min_score=args.min_score, img_size=args.img_size)
    docker run -v $PWD:/opt/app \
    -e PYTHONPATH=$PYTHONPATH:/opt/app \
    -it colemurray/age-gender-estimation-tutorial \
    python3 /opt/app/bin/ \
    --db-path /opt/app/data/imdb_crop/imdb.mat \
    --photo-dir /opt/app/data/imdb_crop \
    --output-dir /opt/app/var \
    --min-score 1.0 \
    --img-size 224

After approximately 20 minutes, you’ll have a processed dataset.

Next, you’ll use TensorFlow’s data pipeline module to provide data to the estimator. is an abstraction to read and manipulate a dataset in parallel, utilizing C++ threads for performance.

Here, you’ll utilize TensorFlow’s CSV Reader to parse the data, preprocess the images, create batches, and shuffle.

import os
import tensorflow as tf
def csv_record_input_fn(img_dir, filenames, img_size=150, repeat_count=-1, shuffle=True,
                        batch_size=16, random=True):
    Creates tensorflow dataset iterator over records from :param{filenames}.
    :param img_dir: Path to directory of cropped images
    :param filenames: array of file paths to load rows from
    :param img_size: size of image
    :param repeat_count: number of times for iterator to repeat
    :param shuffle: flag for shuffling dataset
    :param batch_size: number of examples in batch
    :param random: flag for random distortion to the image
    :return: Iterator of dataset
    def parse_csv_row(line):
        defaults = [[""], [0], [0]]
        filename, age, gender = tf.decode_csv(line, defaults)
        filename = os.path.join(img_dir) + '/' + filename
        image_string = tf.read_file(filename)
        image = tf.image.decode_image(image_string, channels=3)
        image = tf.cast(image, tf.float32)
        image = tf.image.per_image_standardization(image)
        image.set_shape([img_size, img_size, 3])
        age = tf.cast(age, tf.int64)
        gender = tf.cast(gender, tf.int64)
        if random:
            image = tf.image.random_flip_left_right(image)
        return {'image': image}, dict(gender=gender, age=age)
    dataset =
    dataset =
    if shuffle:
        dataset = dataset.shuffle(buffer_size=2000)
    dataset = dataset.batch(batch_size)
    dataset = dataset.repeat(repeat_count)
    dataset = dataset.prefetch(batch_size * 10)
    iterator = dataset.make_one_shot_iterator()
    return iterator.get_next()


Below, you’ll create a basic CNN model. The model consists of three convolutions and two fully connected layers, with a softmax classifier head for each task.

import tensorflow as tf
def network(feature_input, labels, mode):
    Creates a simple multi-layer convolutional neural network
    :param feature_input:
    :param labels:
    :param mode:
    filters = [32, 64, 128]
    dropout_rates = [0.2, 0.4, 0.7]
    conv_layer = feature_input
    for filter_num, dropout_rate in zip(filters, dropout_rates):
        conv_layer = conv_block(conv_layer, mode, filters=filter_num, dropout=dropout_rate)
    # Dense Layer
    pool4_flat = tf.layers.flatten(conv_layer)
    dense = tf.layers.dense(inputs=pool4_flat, units=1024, activation=tf.nn.relu)
    dropout = tf.layers.dropout(
        inputs=dense, rate=0.4, training=mode == tf.estimator.ModeKeys.TRAIN)
    # Age Head
    age_dense = tf.layers.dense(inputs=dropout, units=1024)
    age_logits = tf.layers.dense(inputs=age_dense, units=101)
    # Gender head
    gender_dense = tf.layers.dense(inputs=dropout, units=1024)
    gender_logits = tf.layers.dense(inputs=gender_dense, units=2)
    return age_logits, gender_logits
def conv_block(input_layer, mode, filters=64, dropout=0.0):
    conv = tf.layers.conv2d(
        kernel_size=[5, 5],
    pool = tf.layers.max_pooling2d(inputs=conv, pool_size=[2, 2], strides=2)
    dropout_layer = tf.layers.dropout(
        inputs=pool, rate=dropout, training=mode == tf.estimator.ModeKeys.TRAIN)
    return dropout_layer

Joint loss function

For the training operation, you’ll use the Adam Optimizer. For a loss function, you’ll average the cross-entropy error of each head, creating a shared loss function between the heads.

age and gender joint loss function

TensorFlow estimator

TensorFlow estimators provide a simple abstraction for graph creation and runtime processing. TensorFlow has specified an interface model_fn, that can be used to create custom estimators.

Below, you’ll take the network created above and create training, eval, and predict. These specifications will be used by TensorFlow’s estimator class to alter the behavior of the graph.

import tensorflow as tf
from age_gender_estimation_tutorial.cnn_model import network
def model_fn(features, labels, mode, params):
    Creates model_fn for Tensorflow estimator. This function takes features and input, and
    is responsible for the creation and processing of the Tensorflow graph for training, prediction and evaluation.
    Expected feature: {'image': image tensor }
    :param features: dictionary of input features
    :param labels: dictionary of ground truth labels
    :param mode: graph mode
    :param params: params to configure model
    :return: Estimator spec dependent on mode
    learning_rate = params['learning_rate']
    image_input = features['image']
    age_logits, logits = network(image_input, labels, mode)
    if mode == tf.estimator.ModeKeys.PREDICT:
        return get_prediction_spec(age_logits, logits)
    joint_loss = get_loss(age_logits, logits, labels)
    if mode == tf.estimator.ModeKeys.TRAIN:
        return get_training_spec(learning_rate, joint_loss)
        return get_eval_spec(logits, age_logits, labels, joint_loss)
def get_prediction_spec(age_logits, logits):
    Creates estimator spec for prediction
    :param age_logits: logits of age task
    :param logits: logits of gender task
    :return: Estimator spec
    predictions = {
        "classes": tf.argmax(input=logits, axis=1),
        "age_class": tf.argmax(input=age_logits, name='age_class', axis=1),
        "age_prob": tf.nn.softmax(age_logits, name='age_prob'),
        "probabilities": tf.nn.softmax(logits, name="softmax_tensor")
    return tf.estimator.EstimatorSpec(mode=tf.estimator.ModeKeys.PREDICT, predictions=predictions)
def get_loss(age_logits, gender_logits, labels):
    Creates joint loss function
    :param age_logits: logits of age
    :param gender_logits: logits of gender task
    :param labels: ground-truth labels of age and gender
    :return: joint loss of age and gender
    gender_loss = tf.losses.sparse_softmax_cross_entropy(labels=labels['gender'], logits=gender_logits)
    age_loss = tf.losses.sparse_softmax_cross_entropy(labels=labels['age'], logits=age_logits)
    joint_loss = gender_loss + age_loss
    return joint_loss
def get_eval_spec(gender_logits, age_logits, labels, loss):
    Creates eval spec for tensorflow estimator
    :param gender_logits: logits of gender task
    :param age_logits: logits of age task
    :param labels: ground truth labels for age and gender
    :param loss: loss op
    :return: Eval estimator spec
    eval_metric_ops = {
        "gender_accuracy": tf.metrics.accuracy(
            labels=labels['gender'], predictions=tf.argmax(gender_logits, axis=1)),
        'age_accuracy': tf.metrics.accuracy(labels=labels['age'], predictions=tf.argmax(age_logits, axis=1)),
        'age_precision': tf.metrics.sparse_precision_at_k(labels=labels['age'],
                                                          predictions=age_logits, k=10)
    return tf.estimator.EstimatorSpec(
        mode=tf.estimator.ModeKeys.EVAL, loss=loss, eval_metric_ops=eval_metric_ops)
def get_training_spec(learning_rate, joint_loss):
    Creates training estimator spec
    :param learning rate for optimizer
    :param joint_loss: loss op
    :return: Training estimator spec
    optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate)
    gender_train_op = optimizer.minimize(
    return tf.estimator.EstimatorSpec(mode=tf.estimator.ModeKeys.TRAIN, loss=joint_loss, train_op=gender_train_op)
def serving_fn():
    receiver_tensor = {
        'image': tf.placeholder(dtype=tf.float32, shape=[None, None, None, 3])
    features = {
        'image': tf.image.resize_images(receiver_tensor['image'], [224, 224])
    return tf.estimator.export.ServingInputReceiver(features, receiver_tensor)


Now that you’ve preprocessed the data and created the model architecture and data pipeline, you’ll begin training the model.

import argparse
import tensorflow as tf
from medium_age_estimation_tutorial.cnn_estimator import model_fn, serving_fn
from medium_age_estimation_tutorial.dataset import csv_record_input_fn
if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--img-size', type=int, default=160)
    parser.add_argument('--num-steps', type=int, default=200000)
    args = parser.parse_args()
    config = tf.estimator.RunConfig(model_dir=args.model_dir,
    estimator = tf.estimator.Estimator(
        model_fn=model_fn, config=config, params={
            'learning_rate': 0.0001
    train_spec = tf.estimator.TrainSpec(
        input_fn=lambda: csv_record_input_fn(args.img_dir, args.train_csv, args.img_size, shuffle=False),
    eval_spec = tf.estimator.EvalSpec(
        input_fn=lambda: csv_record_input_fn(args.img_dir, args.val_csv, args.img_size, batch_size=1, shuffle=False,
    tf.estimator.train_and_evaluate(estimator, train_spec, eval_spec)

    docker run -v $PWD:/opt/app \
    -e PYTHONPATH=$PYTHONPATH:/opt/app \
    -it colemurray/age-gender-estimation-tutorial:gpu \
    python3 /opt/app/bin/ \
    --img-dir /opt/app/var/crop \
    --train-csv /opt/app/var/train.csv \
    --val-csv /opt/app/var/val.csv \
    --model-dir /opt/app/var/cnn-model \
    --img-size 224 \
    --num-steps 200000


Below, you’ll load your age and gender TensorFlow model. The model will be loaded from disk and predict on the provided image.

import logging
from argparse import ArgumentParser
import tensorflow as tf
from scipy.misc import imread
from tensorflow.contrib import predictor
logger = logging.getLogger(__name__)
if __name__ == '__main__':
    parser = ArgumentParser(add_help=True)
    parser.add_argument('--model-dir', required=True)
    parser.add_argument('--image-path', required=True)
    args = parser.parse_args()
    prediction_fn = predictor.from_saved_model(export_dir=args.model_dir, signature_def_key='serving_default')
    batch = []
    image = imread(args.image_path)
    output = prediction_fn({
        'image': [image]

    # Update the model path below with your model
    docker run -v $PWD:/opt/app \
    -e PYTHONPATH=$PYTHONPATH:/opt/app \
    -it colemurray/age-gender-estimation-tutorial \
    python3 /opt/app/bin/ \
    --image-path /opt/app/var/crop/25/nm0000325_rm2755562752_1956-1-7_2002.jpg \
    --model-dir /opt/app/var/cnn-model-3/serving/<TIMESTAMP>
Predicted: M/46 Actual: M/46


In this tutorial, you learned how to build and train a multi-task network for predicting a subject’s age and image. By using a shared architecture, both targets can be trained and predicted simultaneously.

Next Steps:

  • Evaluate on Your Own Dataset
  • Try a different network architecture
  • Experiment with Different Hyperparameters

Questions/issues? Open an issue here on GitHub

Complete code here.

Call to Action

If you enjoyed this tutorial, follow and recommend!

Interested in learning more about Deep Learning / Machine Learning? Check out my other tutorials:

- Building an image caption generator with Deep Learning in Tensorflow - Building a Facial Recognition Pipeline with Deep Learning in Tensorflow - Deep Learning CNN’s in Tensorflow with GPUs - Deep Learning with Keras on Google Compute Engine - Recommendation Systems with Apache Spark on Google Compute Engine

Other places you can find me:

Cole Murray (@_ColeMurray) | Twitter