Source code for mxnet.gluon.contrib.estimator.batch_processor

# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

# coding: utf-8
# pylint: disable=wildcard-import, unused-argument, too-many-ancestors
"""Gluon Batch Processor for Estimators"""

from ...utils import split_and_load
from .... import autograd

__all__ = ['BatchProcessor']

[docs]class BatchProcessor(object): """BatchProcessor Class for plug and play fit_batch & evaluate_batch During training or validation, data are divided into minibatches for processing. This class aims at providing hooks of training or validating on a minibatch of data. Users may provide customized fit_batch() and evaluate_batch() methods by inheriting from this class and overriding class methods. :py:class:`BatchProcessor` can be used to replace fit_batch() and evaluate_batch() in the base estimator class """ def __init__(self): pass def _get_data_and_label(self, batch, ctx, batch_axis=0): data = batch[0] label = batch[1] data = split_and_load(data, ctx_list=ctx, batch_axis=batch_axis) label = split_and_load(label, ctx_list=ctx, batch_axis=batch_axis) return data, label
[docs] def evaluate_batch(self, estimator, val_batch, batch_axis=0): """Evaluate the estimator model on a batch of validation data. Parameters ---------- estimator : Estimator Reference to the estimator val_batch : tuple Data and label of a batch from the validation data loader. batch_axis : int, default 0 Batch axis to split the validation data into devices. """ data, label = self._get_data_and_label(val_batch, estimator.context, batch_axis) pred = [estimator.val_net(x) for x in data] loss = [estimator.val_loss(y_hat, y) for y_hat, y in zip(pred, label)] return data, label, pred, loss
[docs] def fit_batch(self, estimator, train_batch, batch_axis=0): """Trains the estimator model on a batch of training data. Parameters ---------- estimator : Estimator Reference to the estimator train_batch : tuple Data and label of a batch from the training data loader. batch_axis : int, default 0 Batch axis to split the training data into devices. Returns ------- data: List of NDArray Sharded data from the batch. Data is sharded with `gluon.split_and_load`. label: List of NDArray Sharded label from the batch. Labels are sharded with `gluon.split_and_load`. pred: List of NDArray Prediction on each of the sharded inputs. loss: List of NDArray Loss on each of the sharded inputs. """ data, label = self._get_data_and_label(train_batch, estimator.context, batch_axis) with autograd.record(): pred = [estimator.net(x) for x in data] loss = [estimator.loss(y_hat, y) for y_hat, y in zip(pred, label)] for l in loss: l.backward() return data, label, pred, loss