PReLU with long compilation time

Cheers,

I have implemented a PReLU nonlinearity that seems to be working quite well in my models, except for the fact the model’s compilation time (*) is way longer than the equivalent model with ReLU instead.

Given my lack of deeper knowledge to go further, may I please ask for your help and advise.

Thanks in advance.

(*) Precompilation prior to running an epoch.

preluweights(ch_in::Int) = DepthwiseConv((1, 1), ch_in => ch_in;
                                 bias=false,
                                 init=rand32
)

struct ConvPReLU
    conv
end
@layer ConvPReLU

function ConvPReLU(ch_in::Int)
    return ConvPReLU(preluweights(ch_in))
end

function (m::ConvPReLU)(x)
    return max.(x, 0) .+ m.conv(min.(x, 0))
end

Update

The below code has been modified, as to implement nonlinearities similarly to what has been done e.g. with ReLU. Execution speed has indeed improved, however precompilation time remains insanely long (such that my ssh terminal breaks the link due to lack of activity).

Any advise is appreciated. Thanks in advance.

# PReLU
preluweights(ch_in::Int) = DepthwiseConv((1, 1), ch_in => ch_in;
                                 bias=false,
                                 init=rand32
)

struct ConvPReLU
    conv
end
@layer ConvPReLU

function ConvPReLU(ch_in::Int)
    return ConvPReLU(preluweights(ch_in))
end

fpos(x) = ifelse(x>0, x, zero(x))   # fpos(x) = max(x, 0)
fneg(x) = ifelse(x<0, x, zero(x))   # fneg(x) = min(x, 0)

function (m::ConvPReLU)(x)
    return fpos.(float(x)) .+ (m.conv(fneg.(float(x))))
end