def get_model(args, num_classes, test=False, tiny=False): """ Create computation graph and variables. Args: tiny: Tiny ImageNet mode if True. """ data_size = 320 nn_in_size = 224 if tiny: data_size = 64 nn_in_size = 56 image = nn.Variable([args.batch_size, 3, data_size, data_size]) label = nn.Variable([args.batch_size, 1]) pimage = image_preprocess(image, nn_in_size, data_size, test) pred, hidden = model_resnet.resnet_imagenet(pimage, num_classes, args.num_layers, args.shortcut_type, test=test, tiny=tiny) loss = F.mean(F.softmax_cross_entropy(pred, label)) Model = namedtuple('Model', ['image', 'label', 'pred', 'loss', 'hidden']) return Model(image, label, pred, loss, hidden)
def get_model(args, num_classes): """ Create computation graph and variables. """ data_size = 224 image = nn.Variable([1, 3, data_size, data_size]) pimage = image_preprocess(image) pred, hidden = model_resnet.resnet_imagenet(pimage, num_classes, args.num_layers, args.shortcut_type, test=True, tiny=False) Model = namedtuple('Model', ['image', 'pred', 'hidden']) return Model(image, pred, hidden)
def get_model(args, num_classes, n_devices, accum_grad, test=False): """ Create computation graph and variables. """ nn_in_size = 224 image = nn.Variable([args.batch_size, 3, nn_in_size, nn_in_size]) label = nn.Variable([args.batch_size, 1]) pred, hidden = model_resnet.resnet_imagenet(image, num_classes, args.num_layers, args.shortcut_type, test=test, tiny=False) loss = F.mean(F.softmax_cross_entropy(pred, label)) / \ (n_devices * accum_grad) Model = namedtuple('Model', ['image', 'label', 'pred', 'loss', 'hidden']) return Model(image, label, pred, loss, hidden)
def get_model(args, num_classes, test=False, tiny=False): """ Create computation graph and variables. Args: tiny: Tiny ImageNet mode if True. """ data_size = 320 nn_in_size = 224 if tiny: data_size = 64 nn_in_size = 56 image = nn.Variable([args.batch_size, 3, data_size, data_size]) label = nn.Variable([args.batch_size, 1]) pimage = image_preprocess(image, nn_in_size) pred, hidden = model_resnet.resnet_imagenet( pimage, num_classes, args.num_layers, args.shortcut_type, test=test, tiny=tiny) loss = F.mean(F.softmax_cross_entropy(pred, label)) Model = namedtuple('Model', ['image', 'label', 'pred', 'loss', 'hidden']) return Model(image, label, pred, loss, hidden)