Added color!
[profiler-torch.git] / profiler.lua
1
2 --[[
3
4    Written by Francois Fleuret (francois@fleuret.org)
5
6    This is free and unencumbered software released into the public
7    domain.
8
9    Anyone is free to copy, modify, publish, use, compile, sell, or
10    distribute this software, either in source code form or as a
11    compiled binary, for any purpose, commercial or non-commercial, and
12    by any means.
13
14    In jurisdictions that recognize copyright laws, the author or
15    authors of this software dedicate any and all copyright interest in
16    the software to the public domain. We make this dedication for the
17    benefit of the public at large and to the detriment of our heirs
18    and successors. We intend this dedication to be an overt act of
19    relinquishment in perpetuity of all present and future rights to
20    this software under copyright law.
21
22    THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
23    EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
24    MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
25    NONINFRINGEMENT.  IN NO EVENT SHALL THE AUTHORS BE LIABLE FOR ANY
26    CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF
27    CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
28    WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
29
30    For more information, please refer to <http://unlicense.org/>
31
32 ]]--
33
34 require 'torch'
35 require 'nn'
36 require 'sys'
37
38 profiler = {}
39
40 profiler.color = true
41
42 profiler.colors = function(name)
43    if profiler.color then
44       return sys.COLORS[name]
45    else
46       return ''
47    end
48 end
49
50 function profiler.decorate(model, functionsToDecorate)
51
52    local functionsToDecorate = functionsToDecorate or
53       {
54          'updateOutput',
55          'backward'
56       }
57
58    for _, name in pairs(functionsToDecorate) do
59       model.accTime = {}
60
61       local nameOrig = name .. '__orig'
62
63       -- We decorate the class and not the object, otherwise we cannot
64       -- save models anymore.
65
66       if rawget(model, name) then
67          error('We decorate the class, not the objects, and there is a ' .. name .. ' in ' .. model)
68       end
69
70       local toDecorate = getmetatable(model)
71
72       if toDecorate[name] and not toDecorate[nameOrig] then
73          toDecorate[nameOrig] = toDecorate[name]
74          toDecorate[name] = function(self, ...)
75             local startTime = sys.clock()
76             local result = { self[nameOrig](self, unpack({...})) }
77             local endTime = sys.clock()
78             self.accTime[name] = (self.accTime[name] or 0) + endTime - startTime
79             return unpack(result)
80          end
81       end
82
83    end
84
85    if torch.isTypeOf(model, nn.Container) then
86       for _, m in ipairs(model.modules) do
87          profiler.decorate(m, functionsToDecorate)
88       end
89    end
90
91 end
92
93 function profiler.timing(l, t, nbSamples, totalTime)
94    local s = string.format('%s %.02fs', l, t)
95    if totalTime then
96       s = s .. string.format(profiler.colors('blue') .. ' [%.02f%%]', 100 * t / totalTime)
97    end
98    if nbSamples then
99       s = s .. string.format(profiler.colors('green') .. ' (%.01fmus/sample)', 1e6 * t / nbSamples)
100    end
101    s = s .. profiler.colors('black')
102    return s
103 end
104
105 function profiler.print(model, nbSamples, totalTime, indent)
106    local indent = indent or ''
107    local hint
108
109    local localTotal = 0
110    for _, t in pairs(model.accTime) do
111       localTotal = localTotal + t
112    end
113
114    if torch.isTypeOf(model, nn.Container) then
115       hint = ' '
116    else
117       if profiler.color then
118          hint = ' '
119       else
120          hint = '*'
121       end
122       hint = hint .. profiler.colors('red')
123    end
124
125    print(profiler.timing(indent .. hint .. ' ' .. model.__typename, localTotal, nbSamples, totalTime))
126
127    for l, t in pairs(model.accTime) do
128       print(profiler.timing(indent .. '  ' .. l, t, nbSamples, totalTime))
129    end
130
131    print()
132
133    if torch.isTypeOf(model, nn.Container) then
134       for _, m in ipairs(model.modules) do
135          profiler.print(m, nbSamples, totalTime, indent .. '  ')
136       end
137    end
138 end
139
140 return profiler