| |
| |
| |
| |
| |
| import unittest |
| import hypothesis.strategies as st |
| from hypothesis import given |
| import numpy as np |
| from caffe2.proto import caffe2_pb2 |
| from caffe2.python import ( |
| brew, |
| core, |
| model_helper, |
| workspace, |
| ) |
| from caffe2.python.transformations import optimizeForMKLDNN |
| import caffe2.python.hypothesis_test_util as hu |
| |
| |
| @unittest.skipIf(not workspace.C.use_mkldnn, "No MKLDNN support.") |
| class PreConvertTest(hu.HypothesisTestCase): |
| @given(input_channels=st.integers(15, 16), |
| batch_size=st.integers(1, 3)) |
| def test_preConvert(self, input_channels, batch_size): |
| def AddModel(model, data): |
| conv1 = brew.conv(model, data, 'conv1', dim_in=input_channels, |
| dim_out=10, kernel=3, stride=1, pad=1, training_mode=1) |
| deconv1 = brew.conv_transpose(model, conv1, 'deconv1', dim_in=10, dim_out=10, |
| kernel=2, stride=2, pad=0, training_mode=1) |
| fc1 = brew.fc(model, deconv1, 'fc1', dim_in=10 * 56 * 56, dim_out=3) |
| softmax = brew.softmax(model, fc1, 'softmax') |
| |
| return softmax |
| |
| def AddTrainingOperators(model, softmax, label): |
| """Adds training operators to the model.""" |
| # Compute cross entropy between softmax scores and labels |
| xent = model.LabelCrossEntropy([softmax, label], 'xent') |
| # Compute the expected loss |
| loss = model.AveragedLoss(xent, "loss") |
| # Use the average loss we just computed to add gradient operators to the model |
| model.AddGradientOperators([loss]) |
| |
| arg_scope = {"order": "NCHW", 'no_bias': False} |
| # Create the model helper for the train model |
| device_opt = core.DeviceOption(caffe2_pb2.IDEEP, 0) |
| with core.DeviceScope(device_opt): |
| train_model = model_helper.ModelHelper(name="test_train", arg_scope=arg_scope) |
| # Add the model definition (fc layers, conv layers, softmax, etc.) |
| softmax = AddModel(train_model, "X") |
| AddTrainingOperators(train_model, softmax, "label") |
| |
| X = np.random.rand( |
| batch_size, input_channels, 28, 28).astype(np.float32) - 0.5 |
| label = np.random.randint(3, size=batch_size).astype(np.int32) |
| blob_dict = {} |
| output_dict = {} |
| output_dict_cosim = {} |
| old_ws_name = workspace.CurrentWorkspace() |
| workspace.FeedBlob('X', X) |
| workspace.FeedBlob('label', label) |
| workspace.RunNetOnce(train_model.param_init_net) |
| for op in train_model.net.Proto().op: |
| if op.type == "Softmax": |
| break |
| for j in range(1, len(op.input)): |
| blob_dict[op.input[j]] = workspace.FetchBlob(op.input[j]) |
| |
| workspace.CreateNet(train_model.net, overwrite=True) |
| optimizeForMKLDNN(train_model.net, training_mode=True) |
| workspace.RunNet(train_model.net) |
| for op in train_model.net.Proto().op: |
| for blob in op.output: |
| output_dict[blob] = workspace.FetchBlob(blob) |
| |
| workspace.SwitchWorkspace("_device_check_", True) |
| workspace.FeedBlob('X', X) |
| workspace.FeedBlob('label', label) |
| for blob in blob_dict.keys(): |
| workspace.FeedBlob(blob, blob_dict[blob]) |
| workspace.CreateNet(train_model.net, overwrite=True) |
| workspace.RunNet(train_model.net) |
| for blob in output_dict.keys(): |
| output_dict_cosim[blob] = workspace.FetchBlob(blob) |
| |
| for blob in output_dict.keys(): |
| if not np.allclose(output_dict[blob], output_dict_cosim[blob], atol=0.001, rtol=0.0001): |
| print("blob {} error".format(blob)) |
| print(np.max(np.abs(output_dict[blob] - output_dict_cosim[blob]))) |
| self.assertTrue(False) |
| |
| workspace.ResetWorkspace() |
| workspace.SwitchWorkspace(old_ws_name) |
| |
| if __name__ == "__main__": |
| unittest.main() |
