├── manifest ├── net-toolkit-scm-1.rockspec ├── test └── test.lua ├── README.md └── init.lua /manifest: -------------------------------------------------------------------------------- 1 | commands = {} 2 | modules = {} 3 | repository = { 4 | ['net-toolkit'] = { 5 | ['scm-1'] = { 6 | { 7 | arch = "rockspec" 8 | } 9 | } 10 | } 11 | } 12 | -------------------------------------------------------------------------------- /net-toolkit-scm-1.rockspec: -------------------------------------------------------------------------------- 1 | package = 'net-toolkit' 2 | version = 'scm-1' 3 | 4 | source = {url = 'git://github.com/Atcold/net-toolkit'} 5 | 6 | description = { 7 | summary = 'A simple module for and the package', 8 | detailed = [[ 9 | It allows to save and retrive to/from disk a lighter version of the network 10 | that is being training. 11 | ]] 12 | } 13 | 14 | dependencies = { 15 | 'torch >= 7.0', 16 | 'nn >= 1.0' 17 | } 18 | 19 | build = { 20 | type = 'builtin', 21 | modules = { 22 | ['net-toolkit.init'] = 'init.lua' 23 | } 24 | } 25 | -------------------------------------------------------------------------------- /test/test.lua: -------------------------------------------------------------------------------- 1 | require 'cunn' 2 | require 'net-toolkit' 3 | 4 | torch.setdefaulttensortype('torch.FloatTensor') 5 | 6 | local batchSize = 64 7 | 8 | -- test CPU 9 | local model = nn.Sequential() 10 | model:add(nn.SpatialConvolutionMM(3, 10, 9, 9)) 11 | model:add(nn.SpatialMaxPooling(2, 2, 2, 2)) 12 | model:add(nn.Threshold(0,0)) 13 | model:add(nn.Dropout(0.5)) 14 | model:add(nn.Reshape(90)) 15 | model:add(nn.Linear(90,10)) 16 | model:add(nn.LogSoftMax()) 17 | 18 | local batch = torch.randn(batchSize, 3, 14, 14) 19 | local gradient = torch.randn(batchSize) 20 | 21 | model:forward(batch) 22 | model:backward(batch, gradient) 23 | 24 | netToolkit.saveNetFields('model-test.net', model, {'weight', 'bias'}) 25 | 26 | local reloadModel = netToolkit.loadNet('model-test.net') 27 | 28 | reloadModel:forward(batch) 29 | reloadModel:backward(batch, gradient) 30 | 31 | os.execute("rm model-test.net") 32 | print('CPU OK') 33 | 34 | -- test GPU 35 | local d_model = nn.Sequential() 36 | d_model:add(nn.Transpose({1,4},{1,3},{1,2})) 37 | d_model:add(nn.SpatialConvolutionCUDA(3, 16, 9, 9)) 38 | d_model:add(nn.SpatialMaxPoolingCUDA(2, 2, 2, 2)) 39 | d_model:add(nn.Threshold(0,0)) 40 | d_model:add(nn.Dropout(0.5)) 41 | d_model:add(nn.Reshape(144)) 42 | d_model:add(nn.Linear(144,10)) 43 | d_model:add(nn.LogSoftMax()) 44 | d_model:cuda() 45 | 46 | local d_batch = torch.CudaTensor(batchSize, 3, 14, 14):copy(batch) 47 | local d_gradient = torch.CudaTensor(batchSize):copy(gradient) 48 | 49 | d_model:forward(d_batch) 50 | d_model:backward(d_batch, d_gradient) 51 | 52 | netToolkit.saveNetFields('model-test.net', d_model, {'weight', 'bias'}) 53 | 54 | local d_reloadModel = netToolkit.loadNet('model-test.net') 55 | 56 | d_reloadModel:forward(d_batch) 57 | d_reloadModel:backward(d_batch, d_gradient) 58 | 59 | os.execute("rm model-test.net") 60 | print('GPU OK') 61 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # net-Toolkit 2 | 3 | A simple module for [`Torch7`](https://github.com/torch/torch7) and the [`nn` package](https://github.com/torch/nn). 4 | 5 | ## Installation 6 | 7 | ```bash 8 | luarocks install --server=https://raw.github.com/Atcold/net-toolkit/master net-toolkit 9 | ``` 10 | 11 | ## Description 12 | 13 | The package `net-toolkit` allows to save and retrive to/from disk a lighter version of the network that is being training. After `require('net-toolkit')`, you will have the following functions at your disposal: 14 | 15 | - [`netToolkit.saveNet()`](#savenet) 16 | - [`netToolkit.saveNetFields()`](#savenetfields) 17 | - [`netToolkit.loadNet()`](#loadnet) 18 | 19 | ### `saveNet()` 20 | 21 | `saveNet()` saves a lighter version of your current network, removing all unnecessary data from it (such as *gradients*, *activation units*' state and etc...) and returns a new couple of flattened *weights* and *gradients*. Usage: 22 | 23 | ```lua 24 | w, dw = saveNet(fileName, model) 25 | ``` 26 | 27 | If you want to save the model in `ascii` format, call `saveNet()` as shown below. The output file will have an `.ascii` extension appended to `fileName`. 28 | 29 | ```lua 30 | w, dw = saveNet(fileName, model, 'ascii') 31 | ``` 32 | 33 | ### `saveNetFields()` 34 | 35 | `saveNetFields()` saves your current network, removing all `Tensor` data you don't want to save and returns a new couple of flattened *weights* and *gradients*. (Set `format` to `'ascii'` for saving in `ascii` format.) Usage: 36 | 37 | ```lua 38 | w, dw = saveNetFields(fileName, model, {'weight', 'bias'}, format) 39 | ``` 40 | 41 | In this case, only `weight` and `bias` `Tensor`s will be saved while the rest will be discarded. (The function `saveNet()` is, therefore, a shortcut to this specific usage of `saveNetFields()`.) However, we could have saved also `gradWeight` and `gradBias` by doing: 42 | 43 | ```lua 44 | w, dw = saveNetFields(fileName, model, {'weight', 'bias', 'gradWeight', 'gradBias'}, format) 45 | ``` 46 | 47 | ### `loadNet()` 48 | 49 | Let's say we would like to load a network we have previously saved with `saveNet()` for continuing a training session on it. Some inner parameters (something about *gradients*) have to be restored, since `saveNet()` did a pruning operation on the network in order to save space. (Set `format` to `'ascii'` if you are trying to load a `.ascii` net. `fileName` does not have to include the `.ascii` extenstion.) Here is how we can handle this case: 50 | 51 | ```lua 52 | model, w, dw = loadNet(fileName, format) 53 | ``` 54 | 55 | Now we can keep training, perhaps without forgetting to (re-)define a [*criterion* `loss`](https://github.com/torch/nn/blob/master/README.md#criterions) (the criterion is not saved with the network, so we have to re-define it, if we don't already do it somewhere else in the code). 56 | -------------------------------------------------------------------------------- /init.lua: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- 2 | -- netToolkit 3 | -------------------------------------------------------------------------------- 4 | -- A simple module for and the package. 5 | -- It allows to save and retrive to/from disk a lighter version of the network 6 | -- that is being training. 7 | -------------------------------------------------------------------------------- 8 | -- Alfredo Canziani, Mar 2014 9 | -------------------------------------------------------------------------------- 10 | 11 | -- Private functions definition ------------------------------------------------ 12 | local function nilling(module, fields) 13 | for key, val in pairs(module) do 14 | if string.match(torch.typename(val) or '', 'Tensor') and not fields[key] then 15 | module[key] = torch.Tensor():typeAs(val) 16 | end 17 | end 18 | end 19 | 20 | local function netLighter(network, fields) 21 | nilling(network, fields) 22 | if network.modules then 23 | for _,a in ipairs(network.modules) do 24 | netLighter(a, fields) 25 | end 26 | end 27 | end 28 | 29 | local function craftGradNBias(module) 30 | if module.weight then module.gradWeight = module.weight:clone() end 31 | if module.bias then module.gradBias = module.bias :clone() end 32 | if module.__typename == 'nn.SpatialConvolutionCUDA' then 33 | module.gradWeightPartial = module.weight:clone() 34 | end 35 | if module.__typename == 'nn.SpatialConvolutionMM' then 36 | module.fgradInput = torch.Tensor():typeAs(module.output) 37 | end 38 | module.gradInput = torch.Tensor():typeAs(module.output) 39 | end 40 | 41 | local function repopulateGradNBias(network) 42 | craftGradNBias(network) 43 | if network.modules then 44 | for _,a in ipairs(network.modules) do 45 | repopulateGradNBias(a) 46 | end 47 | end 48 | end 49 | 50 | -- Public functions definition ------------------------------------------------- 51 | local function saveNetFields(fileName, model, fields, format) 52 | -- Reverse dictionary 53 | local keepFields = {} 54 | for _, val in pairs(fields) do 55 | keepFields[val] = true 56 | end 57 | -- Getting rid of unnecessary things and freeing the memory 58 | netLighter(model, keepFields) 59 | collectgarbage() 60 | fileName = format == 'ascii' and fileName .. '.ascii' or fileName 61 | torch.save(fileName, model, format) 62 | -- Repopulate the gradWeight through the whole net 63 | repopulateGradNBias(model) 64 | -- Return NEW storage for and 65 | return model:getParameters() 66 | end 67 | 68 | local function saveNet(fileName, model, format) 69 | return saveNetFields(fileName, model, {'weight', 'bias'}, format) 70 | end 71 | 72 | local function loadNet(fileName, format) 73 | fileName = format == 'ascii' and fileName .. '.ascii' or fileName 74 | local model = torch.load(fileName, format) 75 | -- Repopulate the gradWeight through the whole net 76 | repopulateGradNBias(model) 77 | return model, model:getParameters() 78 | end 79 | 80 | -- Selecting public functions -------------------------------------------------- 81 | netToolkit = { 82 | saveNet = saveNet, 83 | saveNetFields = saveNetFields, 84 | loadNet = loadNet 85 | } 86 | 87 | -- Returning handle 88 | return netToolkit 89 | --------------------------------------------------------------------------------