Source code for nnabla.models.imagenet.vgg

# Copyright (c) 2017 Sony Corporation. All Rights Reserved.
# Licensed 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
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import absolute_import
import nnabla as nn
from nnabla.utils.nnp_graph import NnpNetworkPass

from nnabla import logger

from .base import ImageNetBase

[docs]class VGG(ImageNetBase): ''' VGG architectures for 11, 13, 16 layers. Args: num_layers (int): Number of layers chosen from 11, 13, 16. The following is a list of string that can be specified to ``use_up_to`` option in ``__call__`` method; * ``'classifier'`` (default): The output of the final affine layer for classification. * ``'pool'``: The output of the final global average pooling. * ``'lastconv'``: The input of the final global average pooling without ReLU activation.. * ``'lastconv+relu'``: Network up to ``'lastconv'`` followed by ReLU activation. References: * `Simonyan and Zisserman, Very Deep Convolutional Networks for Large-Scale Image Recognition. <>`_ ''' def __init__(self, num_layers=11): # Check validity of num_layers set_num_layers = set((11, 13, 16)) assert num_layers in set_num_layers, "num_layers must be chosen from {}".format( set_num_layers) self.num_layers = num_layers # Load nnp self._load_nnp('VGG-{}.nnp'.format(num_layers), 'VGG-{0}/VGG-{0}.nnp'.format(num_layers)) self._KEY_VARIABLE = { 'classifier': 'VGG{}/Affine_3'.format(num_layers), 'pool': 'VGG{}/MaxPooling_5'.format(num_layers), 'lastconv': 'VGG{}/Convolution_12'.format(num_layers), 'lastconv+relu': 'VGG{}/ReLU_12'.format(num_layers), } def _input_shape(self): return (3, 224, 224) def __call__(self, input_var=None, use_from=None, use_up_to='classifier', training=False, returns_net=False, verbose=0): assert use_from is None, 'This should not be set because it is for forward compatibility.' input_var = self.get_input_var(input_var) callback = NnpNetworkPass(verbose) callback.remove_and_rewire('ImageAugmentationX') callback.set_variable('TrainingInput', input_var) callback.set_batch_normalization_batch_stat_all(training) self.use_up_to(use_up_to, callback) if not training: callback.remove_and_rewire( 'VGG{}/Dropout_1'.format(self.num_layers)) callback.remove_and_rewire( 'VGG{}/Dropout_2'.format(self.num_layers)) callback.fix_parameters() batch_size = input_var.shape[0] net = self.nnp.get_network( 'Training', batch_size=batch_size, callback=callback) if returns_net: return net return list(net.outputs.values())[0]