pytorch

1
import numpy as np
2
import unittest
3

4
from caffe2.python import core, workspace, test_util
5

6

7
class TestToyRegression(test_util.TestCase):
8
    def testToyRegression(self):
9
        """Tests a toy regression end to end.
10

11
        The test code carries a simple toy regression in the form
12
            y = 2.0 x1 + 1.5 x2 + 0.5
13
        by randomly generating gaussian inputs and calculating the ground
14
        truth outputs in the net as well. It uses a standard SGD to then
15
        train the parameters.
16
        """
17
        workspace.ResetWorkspace()
18
        init_net = core.Net("init")
19
        W = init_net.UniformFill([], "W", shape=[1, 2], min=-1., max=1.)
20
        B = init_net.ConstantFill([], "B", shape=[1], value=0.0)
21
        W_gt = init_net.GivenTensorFill(
22
            [], "W_gt", shape=[1, 2], values=[2.0, 1.5])
23
        B_gt = init_net.GivenTensorFill([], "B_gt", shape=[1], values=[0.5])
24
        LR = init_net.ConstantFill([], "LR", shape=[1], value=-0.1)
25
        ONE = init_net.ConstantFill([], "ONE", shape=[1], value=1.)
26
        ITER = init_net.ConstantFill([], "ITER", shape=[1], value=0,
27
                                     dtype=core.DataType.INT64)
28

29
        train_net = core.Net("train")
30
        X = train_net.GaussianFill([], "X", shape=[64, 2], mean=0.0, std=1.0)
31
        Y_gt = X.FC([W_gt, B_gt], "Y_gt")
32
        Y_pred = X.FC([W, B], "Y_pred")
33
        dist = train_net.SquaredL2Distance([Y_gt, Y_pred], "dist")
34
        loss = dist.AveragedLoss([], ["loss"])
35
        # Get gradients for all the computations above. Note that in fact we
36
        # don't need to get the gradient the Y_gt computation, but we'll just
37
        # leave it there. In many cases, I am expecting one to load X and Y
38
        # from the disk, so there is really no operator that will calculate the
39
        # Y_gt input.
40
        input_to_grad = train_net.AddGradientOperators([loss], skip=2)
41
        # updates
42
        train_net.Iter(ITER, ITER)
43
        train_net.LearningRate(ITER, "LR", base_lr=-0.1,
44
                               policy="step", stepsize=20, gamma=0.9)
45
        train_net.WeightedSum([W, ONE, input_to_grad[str(W)], LR], W)
46
        train_net.WeightedSum([B, ONE, input_to_grad[str(B)], LR], B)
47
        for blob in [loss, W, B]:
48
            train_net.Print(blob, [])
49

50
        # the CPU part.
51
        plan = core.Plan("toy_regression")
52
        plan.AddStep(core.ExecutionStep("init", init_net))
53
        plan.AddStep(core.ExecutionStep("train", train_net, 200))
54

55
        workspace.RunPlan(plan)
56
        W_result = workspace.FetchBlob("W")
57
        B_result = workspace.FetchBlob("B")
58
        np.testing.assert_array_almost_equal(W_result, [[2.0, 1.5]], decimal=2)
59
        np.testing.assert_array_almost_equal(B_result, [0.5], decimal=2)
60
        workspace.ResetWorkspace()
61

62

63
if __name__ == '__main__':
64
    unittest.main()
65