lenet-fine-tune.lua

-- return function that returns network definition
return function(params)
    -- get original number of classes (10 i.e. one per digit)
    local nclasses = 10

    -- get number of channels from external parameters
    local channels = 1
    -- params.inputShape may be nil during visualization
    if params.inputShape then
        channels = params.inputShape[1]
        assert(params.inputShape[2]==28 and params.inputShape[3]==28, 'Network expects 28x28 images')
    end

    if pcall(function() require('cudnn') end) then
       print('Using CuDNN backend')
       backend = cudnn
       convLayer = cudnn.SpatialConvolution
       convLayerName = 'cudnn.SpatialConvolution'
    else
       print('Failed to load cudnn backend (is libcudnn.so in your library path?)')
       if pcall(function() require('cunn') end) then
           print('Falling back to legacy cunn backend')
       else
           print('Failed to load cunn backend (is CUDA installed?)')
           print('Falling back to legacy nn backend')
       end
       backend = nn -- works with cunn or nn
       convLayer = nn.SpatialConvolutionMM
       convLayerName = 'nn.SpatialConvolutionMM'
    end

    -- -- This is a LeNet model. For more information: http://yann.lecun.com/exdb/lenet/

    local lenet = nn.Sequential()
    lenet:add(nn.MulConstant(0.00390625))
    lenet:add(backend.SpatialConvolution(channels,20,5,5,1,1,0)) -- channels*28*28 -> 20*24*24
    lenet:add(backend.SpatialMaxPooling(2, 2, 2, 2)) -- 20*24*24 -> 20*12*12
    lenet:add(backend.SpatialConvolution(20,50,5,5,1,1,0)) -- 20*12*12 -> 50*8*8
    lenet:add(backend.SpatialMaxPooling(2,2,2,2)) --  50*8*8 -> 50*4*4
    lenet:add(nn.View(-1):setNumInputDims(3))  -- 50*4*4 -> 800
    lenet:add(nn.Linear(800,500))  -- 800 -> 500
    lenet:add(backend.ReLU())
    lenet:add(nn.Linear(500, nclasses))  -- 500 -> nclasses
    lenet:add(nn.LogSoftMax())

    -- multi-GPU implementation needed
    assert(params.ngpus <= 1, "Multi-GPU implementation needed")
    local model = lenet

    local function lenetMnistOddOrEvenFineTune(net)
        -- fix weights of existing layers
        local function dummyAccGradParameters() end
        net:get(2).accGradParameters = dummyAccGradParameters
        net:get(4).accGradParameters = dummyAccGradParameters
        net:get(7).accGradParameters = dummyAccGradParameters
        -- insert 10->2 linear layer
        local l = nn.Linear(10, 2)
        net:insert(l, 10)
        return net
    end

    return {
        model = model,
        loss = nn.ClassNLLCriterion(),
        trainBatchSize = 1,
        validationBatchSize = 100,
        fineTuneHook =lenetMnistOddOrEvenFineTune,
    }
end