| |
|
| |
|
| |
|
| |
|
| |
|
| | import numpy as np |
| | import hypothesis.strategies as st |
| | import unittest |
| | import caffe2.python.hypothesis_test_util as hu |
| | from caffe2.python import core, workspace |
| | from hypothesis import given |
| | import caffe2.python.ideep_test_util as mu |
| |
|
| |
|
| | @unittest.skipIf(not workspace.C.use_mkldnn, "No MKLDNN support.") |
| | class TestMomentumSGDUpdateOps(hu.HypothesisTestCase): |
| | @given(n=st.integers(4, 8), nesterov=st.booleans(), |
| | **mu.gcs) |
| | def test_MomentumSGDUpdate(self, n, nesterov, gc, dc): |
| | param = np.random.rand(n).astype(np.float32) |
| | grad = np.random.rand(n).astype(np.float32) |
| | lr = np.random.rand(1).astype(np.float32) |
| | param_momentum = np.random.rand(n).astype(np.float32) |
| | momentum = 0.9 |
| | op = core.CreateOperator( |
| | "MomentumSGDUpdate", |
| | ["grad", "param_momentum", "lr", "param"], |
| | ["grad", "param_momentum", "param"], |
| | momentum=momentum, |
| | nesterov=int(nesterov), |
| | ) |
| | |
| | input_device_options = {'lr': hu.cpu_do} |
| |
|
| | self.assertDeviceChecks( |
| | dc, |
| | op, |
| | [grad, param_momentum, lr, param], |
| | [0], |
| | input_device_options=input_device_options, |
| | threshold=0.001) |
| |
|
| | op_noparam = core.CreateOperator( |
| | "MomentumSGD", |
| | ["grad", "param_momentum", "lr"], |
| | ["grad", "param_momentum"], |
| | momentum=momentum, |
| | nesterov=int(nesterov), |
| | ) |
| |
|
| | self.assertDeviceChecks( |
| | dc, |
| | op_noparam, |
| | [grad, param_momentum, lr], |
| | [0], |
| | input_device_options=input_device_options, |
| | threshold=0.001) |
| |
|
| |
|
| | if __name__ == "__main__": |
| | unittest.main() |
| |
|
