Search code examples
luadeep-learningtorchsoftmaxresnet

Adding Softmax to ResNet model in Torch


Context: I'm trying to modify this Facebook's ResNet feature extractor script to classify an image and print the ImageNet class label. Let's say I have the model in torch:

local model = torch.load('resnet-101.t7')
local output = model:forward(img:cuda()):squeeze(1)

That gives me the scores for each class. I want to obtain the top 5 classes and their probabilities. I think that to transform scores to probabilities I should use a SoftMax layer first.

So I do:

local model = torch.load('resnet-101.t7')
local softMaxLayer = cudnn.LogSoftMax()
model:add(softMaxLayer)
local output = model:forward(img:cuda()):squeeze(1)

But when I run it I get:

/SpatialSoftMax.lua:38: bad argument #1 to 'resizeAs' (torch.DoubleTensor expected, got torch.CudaTensor)

The model looks good to me: (showing last layers only)

  ...
  (9): cudnn.SpatialAveragePooling(7,7,1,1)
  (10): nn.View(2048)
  (11): nn.Linear(2048 -> 1000)
  (12): cudnn.LogSoftMax
}

Any ideas on what might be wrong?


Solution

  • Layers have types associated with them. By default you are given a double type

    local softMaxLayer = cudnn.LogSoftMax():cuda()
    model:add(softMaxLayer)