mirror of
https://github.com/saymrwulf/pytorch.git
synced 2026-05-15 21:00:47 +00:00
e13101e069
Summary: For MKL-DNN,the filter data will be reorderd to primitive format, it takes a lot of time. So the patch provide a method to convert filter format before training. And "OptimizeForIdeep" will be changed to "OptimizeForMkldnn" in this patch. This patch depends on https://github.com/pytorch/pytorch/pull/12866 Pull Request resolved: https://github.com/pytorch/pytorch/pull/15171 Differential Revision: D14590741 Pulled By: yinghai fbshipit-source-id: 07971c9977edac3c8eec08ca2c39cda639683492
97 lines
4.2 KiB
Python
97 lines
4.2 KiB
Python
from __future__ import absolute_import
|
|
from __future__ import division
|
|
from __future__ import print_function
|
|
from __future__ import unicode_literals
|
|
|
|
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()
|