Caffe2 - Python API
A deep learning, cross platform ML framework
sparse_lengths_sum_benchmark.py
1 # Copyright (c) 2016-present, Facebook, Inc.
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 # http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 ##############################################################################
15 
16 from __future__ import absolute_import
17 from __future__ import division
18 from __future__ import print_function
19 from __future__ import unicode_literals
20 
21 import argparse
22 import numpy as np
23 import datetime
24 
25 from caffe2.python import core, workspace
26 
27 DTYPES = {
28  'uint8': np.uint8,
29  'uint8_fused': np.uint8,
30  'float': np.float32,
31  'float16': np.float16,
32 }
33 
34 
35 def benchmark_sparse_lengths_sum(
36  dtype_str,
37  categorical_limit,
38  embedding_size,
39  average_len,
40  batch_size,
41  iterations):
42  print('Preparing lookup table. ' + str(datetime.datetime.now()))
43 
44  # We will use a constant, but non-trivial value so we save initialization
45  # time.
46  data = np.ones([categorical_limit, embedding_size], dtype=np.float32)
47  data *= 17.01
48 
49  if dtype_str == 'uint8':
50  scale_bias = np.random.rand(categorical_limit, 2).astype(np.float32)
51  workspace.FeedBlob("scale_bias", scale_bias.astype(np.float32))
52  elif dtype_str == 'uint8_fused':
53  scale_bias = np.random.randint(255, size=(categorical_limit, 8))
54  data = np.concatenate([data, scale_bias], axis=1)
55 
56  print('Data has shape {} {}'.format(data.shape, datetime.datetime.now()))
57  workspace.FeedBlob("X", data.astype(DTYPES[dtype_str]))
58 
59  # In order to produce truly random lengths and indices, we will embed a
60  # Python operator in the net to generate them.
61  def f(_, outputs):
62  lengths = np.random.randint(
63  int(average_len * 0.75),
64  int(average_len * 1.25),
65  batch_size).astype(np.int32)
66  indices = np.random.randint(
67  0, categorical_limit, np.sum(lengths)).astype(np.int64)
68  outputs[0].feed(indices)
69  outputs[1].feed(lengths)
70 
71  net = core.Net("mynet")
72  net.Python(f)([], ["indices", "lengths", ])
73  if dtype_str == "uint8":
74  net.SparseLengthsSum8BitsRowwise(["X", "indices", "lengths", "scale_bias"], "Y")
75  elif dtype_str == "uint8_fused":
76  net.SparseLengthsSumFused8BitRowwise(["X", "indices", "lengths"], "Y")
77  else:
78  net.SparseLengthsSum(["X", "indices", "lengths"], "Y")
79  workspace.CreateNet(net)
80 
81  # Set random seed, so that repeated runs will keep the same sequence of
82  # random indices.
83  np.random.seed(1701)
84 
85  print('Preparation finished. ' + str(datetime.datetime.now()))
86 
87  workspace.BenchmarkNet(net.Name(), 1, iterations, True)
88 
89 
90 if __name__ == "__main__":
91  parser = argparse.ArgumentParser(
92  description="minimal benchmark for sparse lengths sum.")
93  parser.add_argument(
94  '-d', "--dtype", choices=list(DTYPES.keys()), default="float",
95  help="The data type for the input lookup table.")
96  parser.add_argument(
97  '-e', "--embedding-size", type=int, default=6000000,
98  help="Lookup table size.")
99  parser.add_argument(
100  "--embedding-dim", type=int, default=128,
101  help="Embedding dimension.")
102  parser.add_argument(
103  "--average_len", type=int, default=27,
104  help="Sparse feature average lengths, default is 27")
105  parser.add_argument(
106  "--batch_size", type=int, default=100,
107  help="The batch size.")
108  parser.add_argument(
109  '-i', "--iteration", type=int, default=100000,
110  help="The number of iterations.")
111  args, extra_args = parser.parse_known_args()
112  core.GlobalInit(['python'] + extra_args)
113  benchmark_sparse_lengths_sum(
114  args.dtype,
115  args.embedding_size,
116  args.embedding_dim,
117  args.average_len,
118  args.batch_size,
119  args.iteration)