pytorch/caffe2/python/optimizer_test.py
Ahmed Taei bcce1bd04a Fix optimizer_context OSS test
Summary:
This will fix the test by querying how many instances of the optimizer are already created.
Because OSS tests doesn't run in isolation causing number of created instances of optimizer to be >= 0.

Reviewed By: akyrola

Differential Revision:
D5462433

Tags: easy

fbshipit-source-id: 7a9ab4fe5345f5d5138abb461ba7a990d9ace840
2017-07-20 12:21:09 -07:00

257 lines
10 KiB
Python

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from caffe2.python.optimizer import (
build_sgd, build_multi_precision_sgd, build_ftrl,
build_adagrad, build_adam, add_weight_decay, SgdOptimizer)
from caffe2.python.optimizer_context import UseOptimizer
from caffe2.python.optimizer_test_util import OptimizerTestBase
from caffe2.python.test_util import TestCase
from caffe2.python import workspace
from caffe2.python.core import DataType
import numpy as np
import unittest
class TestSgd(OptimizerTestBase, TestCase):
def build_optimizer(self, model):
self._skip_gpu = False
return build_sgd(model, base_learning_rate=0.1)
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 TestMultiPrecisionSgd(OptimizerTestBase, TestCase):
def build_optimizer(self, model):
self._skip_gpu = False
return build_multi_precision_sgd(model, base_learning_rate=0.1)
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)
@unittest.skipIf(not workspace.has_gpu_support, "No GPU support")
def testGPUDense(self):
super(TestMultiPrecisionSgd, self).testGPUDense(DataType.FLOAT16)
class TestFtrl(OptimizerTestBase, TestCase):
def build_optimizer(self, model):
self._skip_gpu = True
return build_ftrl(
model, engine=None, alpha=1.0, beta=0.1, lambda1=0.0, lambda2=0.0)
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 TestAdagrad(OptimizerTestBase, TestCase):
def build_optimizer(self, model):
self._skip_gpu = False
return build_adagrad(model, base_learning_rate=1.0)
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 TestAdam(OptimizerTestBase, TestCase):
def build_optimizer(self, model):
self._skip_gpu = False
return build_adam(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)
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.assertEqual(
val,
expected_learning_rate[op.output[0]]
)
self.assertEqual(
expected_weight_grad,
set(),
"Not all weights were decayed: {}".format(expected_weight_grad)
)