Exemplo n.º 1
0
def build_model(inputs, num_classes, is_training, hparams):
    """Constructs the vision model being trained/evaled.

    Args:
      inputs: input features/images being fed to the image model build built.
      num_classes: number of output classes being predicted.
      is_training: is the model training or not.
      hparams: additional hyperparameters associated with the image model.

    Returns:
      The logits of the image model.
    """
    scopes = setup_arg_scopes(is_training)
    with contextlib.nested(*scopes):
        if hparams.model_name == 'pyramid_net':
            logits = build_shake_drop_model(inputs, num_classes, is_training)
        elif hparams.model_name == 'wrn':
            logits = build_wrn_model(inputs, num_classes, hparams.wrn_size)
        elif hparams.model_name == 'shake_shake':
            logits = build_shake_shake_model(inputs, num_classes, hparams,
                                             is_training)
        elif hparams.model_name == 'resnet':
            logits = build_resnet_model(inputs, num_classes, hparams,
                                        is_training)
        else:
            raise ValueError("Unknown model name.")
    return logits
Exemplo n.º 2
0
def build_model(inputs, num_classes, is_training, update_bn, hparams):
    """Constructs the vision model being trained/evaled.

  Args:
    inputs: input features/images being fed to the image model build built.
    num_classes: number of output classes being predicted.
    is_training: is the model training or not.
    hparams: additional hyperparameters associated with the image model.

  Returns:
    The logits of the image model.
  """
    scopes = setup_arg_scopes(is_training)
    with contextlib.nested(*scopes):
        if hparams.model_name == "pyramid_net":
            logits = build_shake_drop_model(inputs, num_classes, is_training)
        elif hparams.model_name == "wrn":
            logits = build_wrn_model(inputs, num_classes, hparams.wrn_size,
                                     update_bn)
        elif hparams.model_name == "shake_shake":
            logits = build_shake_shake_model(inputs, num_classes, hparams,
                                             is_training)
    return logits