FM / python

/usr/local/lib64/python3.6/site-packages/caffe2/python/

/usr/local/lib64/python3.6/site-packages/caffe2/python

Name	Size	Mode	Actions
docs/	-	0755	rm
examples/	-	0755	rm
fakelowp/	-	0755	rm
helpers/	-	0755	rm
ideep/	-	0755	rm
layers/	-	0755	rm
mint/	-	0755	rm
mkl/	-	0755	rm
modeling/	-	0755	rm
models/	-	0755	rm
onnx/	-	0755	rm
operator_test/	-	0755	rm
predictor/	-	0755	rm
rnn/	-	0755	rm
serialized_test/	-	0755	rm
test/	-	0755	rm
trt/	-	0755	rm
__pycache__/	-	0755	rm
allcompare_test.py	2255	0644	edit dl rm
attention.py	12359	0644	edit dl rm
benchmark_generator.py	4912	0644	edit dl rm
binarysize.py	5521	0644	edit dl rm
brew.py	4762	0644	edit dl rm
brew_test.py	11739	0644	edit dl rm
build.py	153	0644	edit dl rm
cached_reader.py	4394	0644	edit dl rm
caffe2_pybind11_state.cpython-36m-x86_64-linux-gnu.so	48299712	0755	edit dl rm
caffe2_pybind11_state_gpu.cpython-36m-x86_64-linux-gnu.so	49048144	0755	edit dl rm
caffe_translator.py	35227	0644	edit dl rm
caffe_translator_test.py	3553	0644	edit dl rm
checkpoint.py	32101	0644	edit dl rm
checkpoint_test.py	13405	0644	edit dl rm
cnn.py	7626	0644	edit dl rm
context.py	2841	0644	edit dl rm
context_test.py	1792	0644	edit dl rm
control.py	19309	0644	edit dl rm
control_ops_grad.py	28893	0644	edit dl rm
control_ops_grad_test.py	1752	0644	edit dl rm
control_ops_util.py	10863	0644	edit dl rm
control_test.py	12276	0644	edit dl rm
convert.py	55	0644	edit dl rm
convert_test.py	201	0644	edit dl rm
convnet_benchmarks.py	20533	0644	edit dl rm
convnet_benchmarks_test.py	839	0644	edit dl rm
core.py	119400	0644	edit dl rm
core_gradients_test.py	38022	0644	edit dl rm
core_test.py	47683	0644	edit dl rm
crf.py	13250	0644	edit dl rm
crf_predict.py	1159	0644	edit dl rm
crf_viterbi_test.py	1663	0644	edit dl rm
dataio.py	23532	0644	edit dl rm
dataio_test.py	17575	0644	edit dl rm
dataset.py	12886	0644	edit dl rm
data_parallel_model.py	83100	0644	edit dl rm
data_parallel_model_test.py	56145	0644	edit dl rm
data_workers.py	15941	0644	edit dl rm
data_workers_test.py	6561	0644	edit dl rm
db_file_reader.py	6608	0644	edit dl rm
db_test.py	1110	0644	edit dl rm
device_checker.py	5157	0644	edit dl rm
dyndep.py	1533	0644	edit dl rm
embedding_generation_benchmark.py	5256	0644	edit dl rm
experiment_util.py	3625	0644	edit dl rm
extension_loader.py	744	0644	edit dl rm
fakefp16_transform_lib.py	322	0644	edit dl rm
filler_test.py	748	0644	edit dl rm
functional.py	4415	0644	edit dl rm
functional_test.py	4204	0644	edit dl rm
fused_8bit_rowwise_conversion_ops_test.py	3945	0644	edit dl rm
gradient_checker.py	15377	0644	edit dl rm
gradient_check_test.py	20729	0644	edit dl rm
gru_cell.py	5129	0644	edit dl rm
hip_test_util.py	405	0644	edit dl rm
hsm_util.py	2259	0644	edit dl rm
hypothesis_test.py	105762	0644	edit dl rm
hypothesis_test_util.py	26853	0644	edit dl rm
ideep_test_util.py	998	0644	edit dl rm
layers_test.py	92931	0644	edit dl rm
layer_model_helper.py	29340	0644	edit dl rm
layer_model_instantiator.py	3935	0644	edit dl rm
layer_parameter_sharing_test.py	9148	0644	edit dl rm
layer_test_util.py	4875	0644	edit dl rm
lazy.py	277	0644	edit dl rm
lazy_dyndep.py	2562	0644	edit dl rm
lazy_dyndep_test.py	3914	0644	edit dl rm
lengths_reducer_fused_8bit_rowwise_ops_test.py	7575	0644	edit dl rm
lengths_reducer_rowwise_8bit_ops_test.py	5710	0644	edit dl rm
lstm_benchmark.py	10649	0644	edit dl rm
memonger.py	34041	0644	edit dl rm
memonger_test.py	36910	0644	edit dl rm
mkl_test_util.py	1142	0644	edit dl rm
model_device_test.py	4777	0644	edit dl rm
model_helper.py	23492	0644	edit dl rm
model_helper_test.py	2336	0644	edit dl rm
modifier_context.py	1772	0644	edit dl rm
muji.py	8131	0644	edit dl rm
muji_test.py	3058	0644	edit dl rm
net_builder.py	27679	0644	edit dl rm
net_builder_test.py	11382	0644	edit dl rm
net_drawer.py	14264	0644	edit dl rm
net_printer.py	12704	0644	edit dl rm
net_printer_test.py	3190	0644	edit dl rm
nomnigraph.py	4216	0644	edit dl rm
nomnigraph_test.py	15427	0644	edit dl rm
nomnigraph_transformations.py	3787	0644	edit dl rm
nomnigraph_transformations_test.py	5767	0644	edit dl rm
normalizer.py	1411	0644	edit dl rm
normalizer_context.py	1007	0644	edit dl rm
normalizer_test.py	487	0644	edit dl rm
numa_benchmark.py	2230	0644	edit dl rm
numa_test.py	1663	0644	edit dl rm
observer_test.py	5316	0644	edit dl rm
operator_fp_exceptions_test.py	1248	0644	edit dl rm
optimizer.py	78813	0644	edit dl rm
optimizer_context.py	1462	0644	edit dl rm
optimizer_test.py	30705	0644	edit dl rm
optimizer_test_util.py	9187	0644	edit dl rm
parallelize_bmuf_distributed_test.py	9908	0644	edit dl rm
parallel_workers.py	7682	0644	edit dl rm
parallel_workers_test.py	3501	0644	edit dl rm
pipeline.py	17283	0644	edit dl rm
pipeline_test.py	2542	0644	edit dl rm
predictor_constants.py	198	0644	edit dl rm
python_op_test.py	9169	0644	edit dl rm
queue_util.py	4459	0644	edit dl rm
record_queue.py	4453	0644	edit dl rm
recurrent.py	13297	0644	edit dl rm
regularizer.py	21120	0644	edit dl rm
regularizer_context.py	1013	0644	edit dl rm
regularizer_test.py	10266	0644	edit dl rm
rnn_cell.py	68233	0644	edit dl rm
schema.py	45621	0644	edit dl rm
schema_test.py	15754	0644	edit dl rm
scope.py	3623	0644	edit dl rm
scope_test.py	5249	0644	edit dl rm
session.py	7642	0644	edit dl rm
session_test.py	2078	0644	edit dl rm
sparse_to_dense_mask_test.py	6565	0644	edit dl rm
sparse_to_dense_test.py	3556	0644	edit dl rm
task.py	24274	0644	edit dl rm
task_test.py	870	0644	edit dl rm
test_util.py	3524	0644	edit dl rm
text_file_reader.py	1990	0644	edit dl rm
timeout_guard.py	4054	0644	edit dl rm
toy_regression_test.py	2822	0644	edit dl rm
transformations.py	1832	0644	edit dl rm
transformations_test.py	11960	0644	edit dl rm
tt_core.py	9349	0644	edit dl rm
tt_core_test.py	2518	0644	edit dl rm
utils.py	14181	0644	edit dl rm
utils_test.py	1399	0644	edit dl rm
visualize.py	6315	0644	edit dl rm
workspace.py	25263	0644	edit dl rm
workspace_test.py	34844	0644	edit dl rm
_import_c_extension.py	2250	0644	edit dl rm
__init__.py	3925	0644	edit dl rm

Edit: /usr/local/lib64/python3.6/site-packages/caffe2/python/optimizer_test.py (30705B)

from caffe2.proto import caffe2_pb2
import caffe2.python.optimizer as optimizer
from caffe2.python.optimizer import (
    build_sgd, build_multi_precision_sgd, build_ftrl, build_gftrl, build_wngrad,
    build_adagrad, build_adadelta, build_adam, build_yellowfin, build_rms_prop,
    build_storm, build_decay_adagrad, add_weight_decay, SgdOptimizer)
from caffe2.python.optimizer_context import UseOptimizer
from caffe2.python.optimizer_test_util import (
    OptimizerTestBase, LRModificationTestBase
)
from caffe2.python import core, workspace
from caffe2.python.test_util import TestCase
import numpy as np
from numpy.testing import assert_allclose, assert_equal
import math
import unittest

class TestLars(OptimizerTestBase, TestCase):
    def testSparse(self):
        raise unittest.SkipTest("no sparse support")

def build_optimizer(self, model, **kwargs):
        self._skip_gpu = False
        return build_sgd(model, base_learning_rate=0.1, lars=0.5, **kwargs)

def check_optimizer(self, optimizer):
        self.assertTrue(optimizer.get_auxiliary_parameters().shared)
        self.assertFalse(optimizer.get_auxiliary_parameters().local)
        for param in optimizer.get_auxiliary_parameters().shared:
            tensor = workspace.FetchBlob(param)
            np.testing.assert_allclose(np.array([1.0]), tensor, atol=1e-5)

class TestMomentumSgd(OptimizerTestBase, TestCase):
    def build_optimizer(self, model, **kwargs):
        self._skip_gpu = False
        return build_sgd(model, base_learning_rate=0.1, momentum=0.1, **kwargs)

def check_optimizer(self, optimizer):
        self.assertTrue(optimizer.get_auxiliary_parameters().shared)
        self.assertTrue(optimizer.get_auxiliary_parameters().local)
        for param in optimizer.get_auxiliary_parameters().shared:
            tensor = workspace.FetchBlob(param)
            np.testing.assert_allclose(np.array([1.0]), tensor, atol=1e-5)

class TestSgd(OptimizerTestBase, LRModificationTestBase, TestCase):
    def build_optimizer(self, model, **kwargs):
        self._skip_gpu = False
        return build_sgd(model, base_learning_rate=0.1, **kwargs)

class TestMultiPrecisionSgd(
    OptimizerTestBase, LRModificationTestBase, TestCase
):
    def build_optimizer(self, model, **kwargs):
        self._skip_gpu = False
        return build_multi_precision_sgd(
            model, base_learning_rate=0.1, **kwargs
        )

@unittest.skipIf(not workspace.has_gpu_support, "No GPU support")
    def testGPUDense(self):
        super(TestMultiPrecisionSgd, self).testGPUDense(core.DataType.FLOAT16)

class TestFtrl(OptimizerTestBase, TestCase):
    def build_optimizer(self, model, **kwargs):
        self._skip_gpu = True
        return build_ftrl(
            model,
            engine=None,
            alpha=1.0,
            beta=0.1,
            lambda1=0.0,
            lambda2=0.0,
            **kwargs
        )

def check_optimizer(self, optimizer):
        self.assertFalse(optimizer.get_auxiliary_parameters().shared)
        self.assertTrue(optimizer.get_auxiliary_parameters().local)
        for param in optimizer.get_auxiliary_parameters().local:
            workspace.FetchBlob(param)

class TestGFtrl(OptimizerTestBase, TestCase):
    def testSparse(self):
        raise unittest.SkipTest("no sparse support")

def build_optimizer(self, model, **kwargs):
        self._skip_gpu = True
        return build_gftrl(
            model,
            engine=None,
            alpha=1.0,
            beta=0.1,
            lambda1=0.0,
            lambda2=0.0,
            **kwargs
        )

class TestAdagrad(OptimizerTestBase, LRModificationTestBase, TestCase):
    def build_optimizer(self, model, **kwargs):
        self._skip_gpu = False
        return build_adagrad(model, base_learning_rate=1.0, lars=0.5, **kwargs)

class TestRowWiseAdagrad(OptimizerTestBase, TestCase):
    def build_optimizer(self, model, **kwargs):
        self._skip_gpu = True
        return build_adagrad(
            model, base_learning_rate=1.0, lars=0.5, rowWise=True, **kwargs
        )

def testDense(self):
        raise unittest.SkipTest("no dense support")

def testGPUDense(self):
        raise unittest.SkipTest("no dense support")

class TestRowWiseAdagradWithCounter(OptimizerTestBase, TestCase):
    def build_optimizer(self, model, **kwargs):
        self._skip_gpu = True
        return build_adagrad(
            model,
            base_learning_rate=1.0,
            lars=0.5,
            rowWise=True,
            counter_halflife=5,
            **kwargs
        )

def check_optimizer(self, optimizer):
        self.assertTrue(optimizer.get_auxiliary_parameters().shared)
        self.assertTrue(optimizer.get_auxiliary_parameters().local)
        self.assertTrue(workspace.HasBlob("optimizer_iteration"))
        iteration_tensor = workspace.FetchBlob("optimizer_iteration")
        np.testing.assert_allclose(np.array([2000]),
                                   iteration_tensor,
                                   atol=1e-5)
        for param in optimizer.get_auxiliary_parameters().shared:
            workspace.FetchBlob(param)
        for param in optimizer.get_auxiliary_parameters().local:
            workspace.FetchBlob(param)

def testDense(self):
        raise unittest.SkipTest("no dense support")

def testGPUDense(self):
        raise unittest.SkipTest("no dense support")

class TestWngrad(OptimizerTestBase, LRModificationTestBase, TestCase):
    def build_optimizer(self, model, **kwargs):
        self._skip_gpu = True
        return build_wngrad(model, base_learning_rate=25.0, **kwargs)

class TestStorm(OptimizerTestBase, LRModificationTestBase, TestCase):
    def build_optimizer(self, model, **kwargs):
        self._skip_gpu = True
        return build_storm(model, base_learning_rate=2.0, **kwargs)

class TestAdadelta(OptimizerTestBase, LRModificationTestBase, TestCase):
    def build_optimizer(self, model, **kwargs):
        self._skip_gpu = False
        return build_adadelta(model, base_learning_rate=1.0, decay=0.995, **kwargs)

class TestAdam(OptimizerTestBase, LRModificationTestBase, TestCase):
    def build_optimizer(self, model, **kwargs):
        self._skip_gpu = False
        return build_adam(model, base_learning_rate=0.1, **kwargs)

def check_optimizer(self, optimizer):
        self.assertTrue(optimizer.get_auxiliary_parameters().shared)
        self.assertTrue(optimizer.get_auxiliary_parameters().local)
        self.assertTrue(workspace.HasBlob("optimizer_iteration"))
        iteration_tensor = workspace.FetchBlob("optimizer_iteration")
        np.testing.assert_allclose(np.array([2000]),
                                   iteration_tensor,
                                   atol=1e-5)
        for param in optimizer.get_auxiliary_parameters().shared:
            workspace.FetchBlob(param)
        for param in optimizer.get_auxiliary_parameters().local:
            workspace.FetchBlob(param)

class TestSmartDecayAdam(OptimizerTestBase, LRModificationTestBase, TestCase):
    def build_optimizer(self, model, **kwargs):
        self._skip_gpu = False
        kwargs['beta1'] = 0.0
        return build_adam(model, base_learning_rate=0.1, use_smart_decay=True, **kwargs)

def check_optimizer(self, optimizer):
        self.assertTrue(optimizer.get_auxiliary_parameters().shared)
        self.assertTrue(optimizer.get_auxiliary_parameters().local)
        self.assertTrue(workspace.HasBlob("optimizer_iteration"))
        blob_names = workspace.Blobs()
        self.assertTrue(any((bn.endswith('_last_seen') for bn in blob_names)))
        for param in optimizer.get_auxiliary_parameters().shared:
            workspace.FetchBlob(param)
        for param in optimizer.get_auxiliary_parameters().local:
            workspace.FetchBlob(param)

class TestDecayAdagrad(OptimizerTestBase, LRModificationTestBase, TestCase):
    def build_optimizer(self, model, **kwargs):
        self._skip_gpu = True
        return build_decay_adagrad(model, base_learning_rate=1.0, **kwargs)

def check_optimizer(self, optimizer):
        self.assertTrue(optimizer.get_auxiliary_parameters().shared)
        self.assertTrue(optimizer.get_auxiliary_parameters().local)
        self.assertTrue(workspace.HasBlob("optimizer_iteration"))
        iteration_tensor = workspace.FetchBlob("optimizer_iteration")
        np.testing.assert_allclose(np.array([2000]),
                                   iteration_tensor,
                                   atol=1e-5)
        for param in optimizer.get_auxiliary_parameters().shared:
            workspace.FetchBlob(param)
        for param in optimizer.get_auxiliary_parameters().local:
            workspace.FetchBlob(param)

def testSparse(self):
        raise unittest.SkipTest("no sparse support")

class TestSparseRAdam(OptimizerTestBase, LRModificationTestBase, TestCase):
    def build_optimizer(self, model, **kwargs):
        self._skip_gpu = True
        return build_adam(model, base_learning_rate=0.1, enableRAdam=True, **kwargs)

def check_optimizer(self, optimizer):
        self.assertTrue(optimizer.get_auxiliary_parameters().shared)
        self.assertTrue(optimizer.get_auxiliary_parameters().local)
        self.assertTrue(workspace.HasBlob("optimizer_iteration"))
        iteration_tensor = workspace.FetchBlob("optimizer_iteration")
        np.testing.assert_allclose(np.array([2000]),
                                   iteration_tensor,
                                   atol=1e-5)
        for param in optimizer.get_auxiliary_parameters().shared:
            workspace.FetchBlob(param)
        for param in optimizer.get_auxiliary_parameters().local:
            workspace.FetchBlob(param)

class TestYellowFin(OptimizerTestBase, TestCase):
    # YellowFin: An automatic tuner for momentum SGD
    # (https://arxiv.org/abs/1706.03471)
    def build_optimizer(self, model):
        self._skip_gpu = False
        return build_yellowfin(model, base_learning_rate=0.1)

def check_optimizer(self, optimizer):
        self.assertTrue(optimizer.get_auxiliary_parameters().shared)
        self.assertTrue(optimizer.get_auxiliary_parameters().local)
        self.assertTrue(workspace.HasBlob("optimizer_iteration"))
        iteration_tensor = workspace.FetchBlob("optimizer_iteration")
        np.testing.assert_allclose(np.array([2000]),
                                   iteration_tensor,
                                   atol=1e-5)
        for param in optimizer.get_auxiliary_parameters().shared:
            workspace.FetchBlob(param)
        for param in optimizer.get_auxiliary_parameters().local:
            workspace.FetchBlob(param)

def testSparse(self):
        raise unittest.SkipTest("no sparse support")

def deb(self, val, beta, i, zero_debias):
        if zero_debias:
            return val / (1.0 - beta ** i)
        else:
            return val

def get_lr_mu(self, distance, grad_var, h_min, h_max):
        # First tune based on dynamic range
        if grad_var == 0:
            dr = h_max / h_min
            mu = ((np.sqrt(dr) - 1) / (np.sqrt(dr) + 1)) ** 2
            lr_min = (1 + np.sqrt(mu)) ** 2 / h_max
            return lr_min, mu

p = distance ** 2 * h_min ** 2 / 2 / grad_var
        w3 = (-math.sqrt(p * p + 4.0 / 27.0 * p * p * p) - p) / 2.0
        w = (1.0 if w3 > 0.0 else -1.0) * math.pow(math.fabs(w3), 1.0 / 3.0)
        y = w - p / 3.0 / w
        root = y + 1
        root = min(root, 1.0 - 1e-6)
        dr = h_max / h_min
        mu = max(((np.sqrt(dr) - 1) / (np.sqrt(dr) + 1)) ** 2, root**2)
        lr_min = (1 - np.sqrt(mu)) ** 2 / h_min
        return lr_min, mu

def caffe2_yellowfin(self, zero_debias, grad_coef, n_dim, n_iter, gpu):
        caffe2_res = {}

alpha = 1.0
        mu = 0.0
        beta = 0.999
        curv_win_width = 20
        epsilon = 1e-6

net = core.Net("net")
        param_init_net = core.Net("param_init_net")
        workspace.ResetWorkspace()

with core.DeviceScope(core.DeviceOption(caffe2_pb2.CPU)):
            iteration = param_init_net.ConstantFill(
                [],
                "iteration",
                shape=[1],
                value=0,
                dtype=core.DataType.INT64)
            iter_mutex = param_init_net.CreateMutex([], ["iteration_mutex"])
            net.AtomicIter([iter_mutex, iteration], [iteration])
        pre_grad = param_init_net.ConstantFill(
            [],
            "pre_grad",
            shape=[n_dim],
            value=grad_coef
        )
        if gpu:
            iteration = net.CopyCPUToGPU(
                [iteration],
                "iteration_cpu"
            )
        iteration_float = net.Cast([iteration], "iteration_float")
        grad = net.Mul([pre_grad, iteration_float], "grad", broadcast=True)
        w = param_init_net.ConstantFill([], "w", shape=[n_dim], value=0.0)

# a hack to create an object with __dict__
        param_info = lambda: None
        param_info.blob = w
        param_info.grad = grad

optimizer.YellowFinOptimizer(
            alpha=alpha,
            mu=mu,
            beta=beta,
            curv_win_width=curv_win_width,
            epsilon=epsilon,
            zero_debias=zero_debias
        )._run(
            net,
            param_init_net,
            param_info
        )

workspace.RunNetOnce(param_init_net)
        workspace.CreateNet(net, overwrite=True)
        for i in range(n_iter):
            workspace.RunNet(net)
            scalars_memory_blob = workspace.FetchBlob("w_scalars_memory")
            g_norm2_avg = scalars_memory_blob[1]
            g_norm2_min_avg = scalars_memory_blob[2]
            g_norm2_max_avg = scalars_memory_blob[3]
            distance_avg = scalars_memory_blob[4]
            g_avg_blob = workspace.FetchBlob("w_g_avg")
            res_lr = workspace.FetchBlob("w_lr_avg")[0]
            res_mu = workspace.FetchBlob("w_mu_avg")[0]
            g_deb = self.deb(g_avg_blob, beta, i + 1, zero_debias)
            variance = max(
                self.deb(g_norm2_avg, beta, i + 1, zero_debias) -
                g_deb.dot(g_deb),
                epsilon
            )
            if i > 0:
                caffe2_res[i] = {
                    'h_max': np.exp(self.deb(g_norm2_max_avg,
                                             beta,
                                             i + 1,
                                             zero_debias)),
                    'h_min': np.exp(self.deb(g_norm2_min_avg,
                                             beta,
                                             i + 1,
                                             zero_debias)),
                    'var': variance,
                    'dist': self.deb(distance_avg, beta, i + 1, zero_debias),
                    'lr': res_lr,
                    'mu': res_mu
                }
        return caffe2_res

def numpy_yellowfin(self, zero_debias, grad_coef, n_dim, n_iter, gpu):
        numpy_res = {}

target_h_max = 0.0
        target_h_min = 0.0
        target_g_norm_squared_avg = 0.0
        target_g_norm_avg = 0.0
        target_g_avg = 0.0
        target_dist_avg = 0.0
        target_lr = 1.0
        target_mu = 0.0

for i in range(n_iter):
            grad_val = (i + 1) * grad_coef
            target_g_norm_squared_avg = 0.999 * target_g_norm_squared_avg + \
                0.001 * np.sum((grad_val * np.ones([n_dim, ])) ** 2)
            target_g_norm_avg = 0.999 * target_g_norm_avg + \
                0.001 * np.linalg.norm(grad_val * np.ones([n_dim, ]))
            target_g_avg = 0.999 * target_g_avg + 0.001 * grad_val

target_h_max = 0.999 * target_h_max + \
                0.001 * np.log(grad_val ** 2 * n_dim)
            target_h_min = 0.999 * target_h_min + \
                0.001 * np.log((max(1, i + 2 - 20) * grad_coef) ** 2 * n_dim)
            if zero_debias:
                target_var = target_g_norm_squared_avg / \
                    (1 - 0.999 ** (i + 1)) - \
                    target_g_avg ** 2 * n_dim / (1 - 0.999 ** (i + 1)) ** 2
            else:
                target_var = target_g_norm_squared_avg - \
                    target_g_avg ** 2 * n_dim
            target_dist_avg = 0.999 * target_dist_avg + \
                0.001 * target_g_norm_avg / target_g_norm_squared_avg

if i > 0:
                if zero_debias:
                    lr, mu = self.get_lr_mu(
                        target_dist_avg / (1.0 - 0.999 ** (i + 1)),
                        target_var,
                        np.exp(target_h_min / (1.0 - 0.999 ** (i + 1))),
                        np.exp(target_h_max / (1.0 - 0.999 ** (i + 1))))
                    target_lr = 0.999 * target_lr + 0.001 * lr
                    target_mu = 0.999 * target_mu + 0.001 * mu
                    numpy_res[i] = {
                        'h_max': np.exp(target_h_max / (1 - 0.999 ** (i + 1))),
                        'h_min': np.exp(target_h_min / (1 - 0.999 ** (i + 1))),
                        'var': target_var,
                        'dist': target_dist_avg / (1 - 0.999 ** (i + 1)),
                        'lr': target_lr,
                        'mu': target_mu
                    }
                else:
                    lr, mu = self.get_lr_mu(
                        target_dist_avg,
                        target_var,
                        np.exp(target_h_min),
                        np.exp(target_h_max))
                    target_lr = 0.999 * target_lr + 0.001 * lr
                    target_mu = 0.999 * target_mu + 0.001 * mu
                    numpy_res[i] = {
                        'h_max': np.exp(target_h_max),
                        'h_min': np.exp(target_h_min),
                        'var': target_var,
                        'dist': target_dist_avg,
                        'lr': target_lr,
                        'mu': target_mu
                    }
        return numpy_res

def compare_yellowfin_models(self,
                                 model0,
                                 model1,
                                 zero_debias,
                                 grad_coef,
                                 n_dim,
                                 n_iter,
                                 gpu):
        model0_res = model0(zero_debias, grad_coef, n_dim, n_iter, gpu)
        model1_res = model1(zero_debias, grad_coef, n_dim, n_iter, gpu)
        assert_equal(len(model0_res), len(model1_res))
        for i in range(1, len(model0_res)):
            assert_equal(model0_res[i].keys(), model1_res[i].keys())
            for feat in model0_res[i].keys():
                err_msg = \
                    'i=' + str(i) + ',\n' + \
                    'feat=' + feat + ',\n' + \
                    'grad_coef=' + str(grad_coef) + ',\n' + \
                    'zero_debias=' + str(zero_debias)
                assert_allclose(model0_res[i][feat],
                                model1_res[i][feat],
                                rtol=1e-2,
                                err_msg=err_msg)

@unittest.skip("Results might vary too much. Only for individual use.")
    def test_caffe2_cpu_vs_numpy(self):
        n_dim = 1000000
        n_iter = 50
        cpu_device_opt = core.DeviceOption(caffe2_pb2.CPU)
        with core.DeviceScope(cpu_device_opt):
            for zero_debias, grad_coef in [
                (False, 1.0),
                (False, 0.1),
                (False, 0.01),
                (True, 1.0)
            ]:
                self.compare_yellowfin_models(
                    self.caffe2_yellowfin,
                    self.numpy_yellowfin,
                    zero_debias,
                    grad_coef,
                    n_dim,
                    n_iter,
                    gpu=False
                )

@unittest.skip("Results might vary too much. Only for individual use.")
    @unittest.skipIf(not workspace.has_gpu_support, "No gpu support")
    def test_caffe2_gpu_vs_numpy(self):
        n_dim = 1000000
        n_iter = 50
        gpu_device_opt = core.DeviceOption(workspace.GpuDeviceType, 0)
        with core.DeviceScope(gpu_device_opt):
            for zero_debias in [False, True]:
                for grad_coef in [1.0, 0.1, 0.01]:
                    self.compare_yellowfin_models(
                        self.caffe2_yellowfin,
                        self.numpy_yellowfin,
                        zero_debias,
                        grad_coef,
                        n_dim,
                        n_iter,
                        gpu=True
                    )

class TestRmsProp(OptimizerTestBase, LRModificationTestBase, TestCase):
    def build_optimizer(self, model, **kwargs):
        self._skip_gpu = False
        return build_rms_prop(
            model, base_learning_rate=0.1, epsilon=0.1, **kwargs
        )

def testSparse(self):
        raise unittest.SkipTest("no sparse support")

class TestMultiOptimizers(TestCase):
    def test_multiple_optimizers(self):
        from caffe2.python import brew, core, optimizer
        from caffe2.python.model_helper import ModelHelper

model = ModelHelper(name="test")
        fc1 = brew.fc(model, 'data', 'fc1', 100, 50)
        fc2 = brew.fc(model, fc1, 'fc2', 50, 25)
        pred = brew.fc(model, fc2, 'fc3', 25, 10)
        (softmax, loss) = model.SoftmaxWithLoss(
            [pred, 'label'],
            ['softmax', 'loss'],
        )
        model.AddGradientOperators([loss])

param_to_device = optimizer._get_param_to_device(model)

def infer_blob_device(blob_name):
            return optimizer.get_param_device(
                blob_name, "{}_grad".format(blob_name), param_to_device
            )

sgd_1 = optimizer.SgdOptimizer(base_learning_rate=0.1)
        sgd_2 = optimizer.SgdOptimizer(base_learning_rate=0.2)
        adagrad = optimizer.AdagradOptimizer()

# Check same optimizer share the same learning rate.
        with core.DeviceScope(infer_blob_device("fc1_w")):
            sgd_1(model.net, model.param_init_net, "fc1_w", "fc1_w_grad")
        with core.DeviceScope(infer_blob_device("fc1_b")):
            sgd_1(model.net, model.param_init_net, "fc1_b", "fc1_b_grad")
        fc1_lr_blobs = []
        for op in model.net.Proto().op:
            if op.type == 'WeightedSum' and op.input[0] == 'fc1_w' or \
                    op.input[0] == 'fc1_b':
                fc1_lr_blobs.append(op.input[3])
        self.assertEqual(fc1_lr_blobs[0], fc1_lr_blobs[1])

# Check different instance of the same optimizer has a different lr.
        with core.DeviceScope(infer_blob_device("fc2_w")):
            sgd_2(model.net, model.param_init_net, "fc2_w", "fc2_w_grad")
        with core.DeviceScope(infer_blob_device("fc2_b")):
            sgd_2(model.net, model.param_init_net, "fc2_b", "fc2_b_grad")
        fc2_lr_blobs = []
        for op in model.net.Proto().op:
            if op.type == 'WeightedSum' and op.input[0] == 'fc2_w' or \
                    op.input[0] == 'fc2_b':
                self.assertTrue(op.input[3] not in fc1_lr_blobs)
                fc2_lr_blobs.append(op.input[3])
        self.assertEqual(fc2_lr_blobs[0], fc2_lr_blobs[1])

# Check different optimizer type case
        with core.DeviceScope(infer_blob_device("fc3_w")):
            adagrad(model.net, model.param_init_net, "fc3_w", "fc3_w_grad")
        with core.DeviceScope(infer_blob_device("fc3_b")):
            adagrad(model.net, model.param_init_net, "fc3_b", "fc3_b_grad")
        fc3_lr_blobs = []
        for op in model.net.Proto().op:
            if op.type == 'Adagrad' and op.input[0] == 'fc3_w' or \
                    op.input[0] == 'fc3_b':
                self.assertTrue(op.input[3] not in fc2_lr_blobs)
                self.assertTrue(op.input[3] not in fc1_lr_blobs)
                fc3_lr_blobs.append(op.input[3])
        self.assertEqual(fc3_lr_blobs[0], fc3_lr_blobs[1])

class TestWeightDecay(TestCase):

def test_weight_decay(self):
        from caffe2.python import brew
        from caffe2.python.model_helper import ModelHelper

model = ModelHelper(name="test", arg_scope={'order': 'NCHW'})
        cnv = brew.conv(model, 'data', 'cnv', 32, 32, 4)
        a = brew.fc(model, cnv, 'a', 100, 200)
        pred = brew.fc(model, a, 'b', 200, 5)
        (softmax, loss) = model.SoftmaxWithLoss(
            [pred, 'label'],
            ['softmax', 'loss'],
        )
        model.AddGradientOperators([loss])

add_weight_decay(model, weight_decay=1e-4)
        build_sgd(model, 0.11)

expected_weight_grad = {'b_w_grad', 'a_w_grad', 'cnv_w_grad'}

# Check the proto that all weights are decayed and not non-weights
        # are decayed.
        for op in model.net.Proto().op:
            if op.type == 'WeightedSum' and 'wd_0_0' in op.input:
                if op.output[0] not in expected_weight_grad:
                    print(
                        "Unexpected param for weight_decay: {}".
                        format(op.output[0])
                    )
                self.assertTrue(op.output[0] in expected_weight_grad)
                expected_weight_grad.remove(op.output[0])

self.assertEqual(
            expected_weight_grad,
            set(),
            "Not all weights were decayed: {}".format(expected_weight_grad)
        )

class TestOptimizerContext(TestCase):

def test_optimizer_context(self):
        from caffe2.python import brew, optimizer
        from caffe2.python.model_helper import ModelHelper

model = ModelHelper(name="test", arg_scope={'order': 'NCHW'})
        count = optimizer._optimizer_instance_count['SgdOptimizer']
        cnv_optim = SgdOptimizer(0.15)
        weight_optim = SgdOptimizer(0.2)
        bias_optim = SgdOptimizer(0.1)

with UseOptimizer(cnv_optim):
            cnv = brew.conv(model, 'data', 'cnv', 32, 32, 4)
        with UseOptimizer({'WEIGHT': weight_optim, 'BIAS': bias_optim}):
            a = brew.fc(model, cnv, 'a', 100, 200)
        pred = brew.fc(model, a, 'b', 200, 5)
        (softmax, loss) = model.SoftmaxWithLoss(
            [pred, 'label'],
            ['softmax', 'loss'],
        )
        model.AddGradientOperators([loss])

add_weight_decay(model, weight_decay=1e-4)
        # use the following optimizer if none specified in param_info
        build_sgd(model, 0.11)
        expected_weight_grad = {'b_w_grad', 'a_w_grad', 'cnv_w_grad'}
        expected_learning_rate = {
            "SgdOptimizer_{}_lr_cpu".format(count): -0.15,
            "SgdOptimizer_{}_lr_cpu".format(count + 1): -0.2,
            "SgdOptimizer_{}_lr_cpu".format(count + 2): -0.1,
            "SgdOptimizer_{}_lr_cpu".format(count + 3): -0.11
        }

for op in model.net.Proto().op:
            # Check the proto that all weights are decayed and not non-weights
            # are decayed.
            if op.type == 'WeightedSum' and 'wd_0_0' in op.input:
                if op.output[0] not in expected_weight_grad:
                    print(
                        "Unexpected param for weight_decay: {}".
                        format(op.output[0])
                    )
                self.assertTrue(op.output[0] in expected_weight_grad)
                expected_weight_grad.remove(op.output[0])
            # Check the learning rate for each parameter
            if op.type == 'LearningRate':
                val = 0
                for arg in op.arg:
                    if arg.name == 'base_lr':
                        val = arg.f
                self.assertAlmostEqual(
                    val,
                    expected_learning_rate[op.output[0]]
                )

self.assertEqual(
            expected_weight_grad,
            set(),
            "Not all weights were decayed: {}".format(expected_weight_grad)
        )

cmd: