Torch：我的自定义层不适用于 GPU

Question

我在torch上实现了一个自定义的softmax层，它在CPU模式下工作得很好，但是当我把它改成GPU模式时，它就不能工作了，我的模型的损失在我做一些前向和后向迭代后保持不变，在 CPU 模式下，它会在每个 epoch 后减小。

我的自定义 softmax 层是

local MySoftMax, Parent = torch.class('nn.MySoftMax', 'nn.Module')
function MySoftMax:__init()
    Parent.__init(self)
end
--input = batch * 2 * 40 * 40 * 40
--output = batch * 1 * 40 * 40 * 40

function MySoftMax:updateOutput(input)
    tmp_input = input:clone()
    size_input = tmp_input:size()
    self.output = torch.CudaTensor(size_input[1], 1, size_input[3], size_input[4], size_input[5])
    exp_input = tmp_input:exp()
    for i = 1, size_input[1] do
        tmp_output = exp_input[i][1]:clone()
        sum_exp_input = exp_input[i][1]:add(exp_input[i][2])
        self.output[i][1] = tmp_output:cdiv(sum_exp_input)
    end
    return self.output
end

function MySoftMax:updateGradInput(input, gradOutput)
    tmp_input = input:clone()
    size_input = input:size()
    self.gradInput = torch.CudaTensor(size_input[1], 2, size_input[3], size_input[4], size_input[5])
    exp_input = tmp_input:exp()
    for i = 1, size_input[1] do
        tmp_exp_input = exp_input[i][1]:clone()
        add_exp = exp_input[i][1]:add(exp_input[i][2])
        pow_exp = add_exp:pow(2)
        mul_exp = tmp_exp_input:cmul(exp_input[i][2])

        self.gradInput[i][1] = mul_exp:cdiv(pow_exp)
        self.gradInput[i][2] = self.gradInput[i][1]
        self.gradInput[i][2] = -self.gradInput[i][1]
        self.gradInput[i][1]:cmul(gradOutput[i][1])
        self.gradInput[i][2]:cmul(gradOutput[i][1])
    end
    return self.gradInput
end

GPU 训练代码为：

model = model:cuda()
criterion = criterion:cuda()
inputs = torch.CudaTensor(batch_size, 1, bbx_size, bbx_size, bbx_size)
targets = torch.CudaTensor(batch_size, 1, bbx_size, bbx_size, bbx_size)
model:training()
torch.manualSeed(train_size)
shuffle = torch.randperm(train_size)
print '==> epoch = '
print(epoch)
epoch_error = 0.0
for t = 1, train_size, batch_size do
    xlua.progress(t, train_size)
    for i = 1, batch_size do 
        inputs[i] = train_input[shuffle[t + i - 1]]
        targets[i] = train_output[shuffle[t + i  - 1]]
    end
    local feval = function(x)
        gradParameters:zero()
        outputs = model:forward(inputs)
        err = criterion:forward(outputs, targets)
        df_do = criterion:backward(outputs, targets)
        model:backward(inputs, df_do)
        epoch_error = epoch_error + err
        return f, gradParameters
    end
    optim.sgd(feval, parameters, config)
end

我的模型是：

cnn_state = {8, 16, 32, 64, 128}
pool_size = 2
filter_size = 2
batch_size = 64

model = nn.Sequential()

--stage 1
--1*40*40*40
model:add(nn.VolumetricConvolution(1, cnn_state[1], filter_size, filter_size, filter_size))
model:add(nn.ReLU())
--8*39*39*39
model:add(nn.VolumetricConvolution(cnn_state[1], cnn_state[2], filter_size, filter_size, filter_size))
model:add(nn.ReLU())
--16*38*38*38
maxpool1 = nn.VolumetricMaxPooling(pool_size, pool_size, pool_size, pool_size, pool_size, pool_size)
model:add(maxpool1)
--16*19*19*19

--stage 2
--16*19*19*19
model:add(nn.VolumetricConvolution(cnn_state[2], cnn_state[3], filter_size, filter_size, filter_size))
model:add(nn.ReLU())
--32*18*18*18
model:add(nn.VolumetricConvolution(cnn_state[3], cnn_state[4], filter_size + 1, filter_size + 1, filter_size + 1))
model:add(nn.ReLU())
--64*16*16*16
maxpool2 = nn.VolumetricMaxPooling(pool_size, pool_size, pool_size, pool_size, pool_size, pool_size)
model:add(maxpool2)
--64*8*8*8

--stage 3
model:add(nn.VolumetricConvolution(cnn_state[4],cnn_state[5], filter_size, filter_size, filter_size))
model:add(nn.ReLU())
--128*7*7*7
--deconvolution
model:add(nn.VolumetricFullConvolution(cnn_state[5], cnn_state[4], filter_size, filter_size, filter_size))
model:add(nn.ReLU())
--64*8*8*8
model:add(nn.VolumetricMaxUnpooling(maxpool2))
--64*16*16*16
model:add(nn.VolumetricFullConvolution(cnn_state[4], cnn_state[3], filter_size + 1, filter_size + 1, filter_size + 1))
model:add(nn.ReLU())
--32*18*18*18
model:add(nn.VolumetricFullConvolution(cnn_state[3], cnn_state[2], filter_size, filter_size, filter_size))
model:add(nn.ReLU())
--16*19*19*19
model:add(nn.VolumetricMaxUnpooling(maxpool1))
--16*38*38*38
model:add(nn.VolumetricFullConvolution(cnn_state[2], cnn_state[1], filter_size, filter_size, filter_size))
model:add(nn.ReLU())
--8*39*39*39
model:add(nn.VolumetricFullConvolution(cnn_state[1], 2, filter_size, filter_size, filter_size))
--2*40*40*40
model:add(nn.MySoftMax())

我的模型在 CPU 模式下运行良好，如果我使用 Tenser 而不是 CudaTensor，它的损失会减少并且预测会变得越来越好。但是在 GPU 模式下，它不起作用。 我认为问题可能在于我的自定义层实现，也许它不适用于 GPU。所以我想知道在 GPU 上实现层运行时有什么需要注意的吗？还是我的代码有什么问题？

Answer 1

我不确定在您的特定情况下会发生这种情况，但通常您不应更改状态变量的指针，例如gradInput。反向传播可能以某种方式依赖于此。尝试更改此行

self.gradInput = torch.CudaTensor(size_input[1], 2, size_input[3], size_input[4], size_input[5])

用这一行：

self.gradInput.copy(torch.CudaTensor(size_input[1], 2, size_input[3], size_input[4], size_input[5]))