Skip to content

Commit

Permalink
Moving closer to the one true function
Browse files Browse the repository at this point in the history
  • Loading branch information
@theabhirath
theabhirath committed Aug 16, 2022
1 parent fc03d70 commit 8e54d47
Show file tree
Hide file tree
Showing 5 changed files with 89 additions and 87 deletions.
52 changes: 28 additions & 24 deletions src/convnets/efficientnets/core.jl
View file
Original file line number Diff line number Diff line change
@@ -1,62 +1,66 @@
function mbconv_builder(block_configs::AbstractVector{NTuple{6, Int}},
stage_idx::Integer; scalings::NTuple{2, Real} = (1, 1),
norm_layer = BatchNorm)
function mbconv_builder(block_configs::AbstractVector{<:Tuple},
inplanes::Integer, stage_idx::Integer;
scalings::NTuple{2, Real} = (1, 1), norm_layer = BatchNorm,
round_fn = planes -> _round_channels(planes, 8))
width_mult, depth_mult = scalings
k, inplanes, outplanes, expansion, stride, nrepeats = block_configs[stage_idx]
inplanes = _round_channels(inplanes * width_mult, 8)
k, outplanes, expansion, stride, nrepeats, reduction, activation = block_configs[stage_idx]
inplanes = stage_idx == 1 ? inplanes : block_configs[stage_idx - 1][2]
inplanes = round_fn(inplanes * width_mult)
outplanes = _round_channels(outplanes * width_mult, 8)
function get_layers(block_idx)
inplanes = block_idx == 1 ? inplanes : outplanes
explanes = _round_channels(inplanes * expansion, 8)
stride = block_idx == 1 ? stride : 1
block = mbconv((k, k), inplanes, explanes, outplanes, swish; norm_layer,
stride, reduction = 4)
block = mbconv((k, k), inplanes, explanes, outplanes, activation; norm_layer,
stride, reduction)
return stride == 1 && inplanes == outplanes ? (identity, block) : (block,)
end
return get_layers, ceil(Int, nrepeats * depth_mult)
end

function fused_mbconv_builder(block_configs::AbstractVector{NTuple{6, Int}},
stage_idx::Integer; scalings::NTuple{2, Real} = (1, 1),
norm_layer = BatchNorm)
k, inplanes, outplanes, expansion, stride, nrepeats = block_configs[stage_idx]
function fused_mbconv_builder(block_configs::AbstractVector{<:Tuple},
inplanes::Integer, stage_idx::Integer;
scalings::NTuple{2, Real} = (1, 1), norm_layer = BatchNorm)
k, outplanes, expansion, stride, nrepeats, _, activation = block_configs[stage_idx]
inplanes = stage_idx == 1 ? inplanes : block_configs[stage_idx - 1][2]
function get_layers(block_idx)
inplanes = block_idx == 1 ? inplanes : outplanes
explanes = _round_channels(inplanes * expansion, 8)
stride = block_idx == 1 ? stride : 1
block = fused_mbconv((k, k), inplanes, explanes, outplanes, swish;
block = fused_mbconv((k, k), inplanes, explanes, outplanes, activation;
norm_layer, stride)
return stride == 1 && inplanes == outplanes ? (identity, block) : (block,)
end
return get_layers, nrepeats
end

function efficientnet_builder(block_configs::AbstractVector{NTuple{6, Int}},
residual_fns::AbstractVector;
scalings::NTuple{2, Real} = (1, 1), norm_layer = BatchNorm)
bxs = [residual_fn(block_configs, stage_idx; scalings, norm_layer)
function mbconv_stack_builder(block_configs::AbstractVector{<:Tuple},
residual_fns::AbstractVector; inplanes::Integer,
scalings::NTuple{2, Real} = (1, 1),
norm_layer = BatchNorm)
bxs = [residual_fn(block_configs, inplanes, stage_idx; scalings, norm_layer)
for (stage_idx, residual_fn) in enumerate(residual_fns)]
return (stage_idx, block_idx) -> first.(bxs)[stage_idx](block_idx), last.(bxs)
end

function efficientnet(block_configs::AbstractVector{NTuple{6, Int}},
residual_fns::AbstractVector; scalings::NTuple{2, Real} = (1, 1),
function efficientnet(block_configs::AbstractVector{<:Tuple},
residual_fns::AbstractVector; inplanes::Integer,
scalings::NTuple{2, Real} = (1, 1),
headplanes::Integer = block_configs[end][3] * 4,
norm_layer = BatchNorm, dropout_rate = nothing,
inchannels::Integer = 3, nclasses::Integer = 1000)
layers = []
# stem of the model
append!(layers,
conv_norm((3, 3), inchannels,
_round_channels(block_configs[1][2] * scalings[1], 8), swish;
norm_layer, stride = 2, pad = SamePad()))
conv_norm((3, 3), inchannels, _round_channels(inplanes * scalings[1], 8),
swish; norm_layer, stride = 2, pad = SamePad()))
# building inverted residual blocks
get_layers, block_repeats = efficientnet_builder(block_configs, residual_fns;
scalings, norm_layer)
get_layers, block_repeats = mbconv_stack_builder(block_configs, residual_fns;
inplanes, scalings, norm_layer)
append!(layers, resnet_stages(get_layers, block_repeats, +))
# building last layers
append!(layers,
conv_norm((1, 1), _round_channels(block_configs[end][3] * scalings[1], 8),
conv_norm((1, 1), _round_channels(block_configs[end][2] * scalings[1], 8),
headplanes, swish; pad = SamePad()))
return Chain(Chain(layers...), create_classifier(headplanes, nclasses; dropout_rate))
end
18 changes: 9 additions & 9 deletions src/convnets/efficientnets/efficientnet.jl
View file
Original file line number Diff line number Diff line change
@@ -1,13 +1,13 @@
# block configs for EfficientNet
const EFFICIENTNET_BLOCK_CONFIGS = [
# k, i, o, e, s, n
(3, 32, 16, 1, 1, 1),
(3, 16, 24, 6, 2, 2),
(5, 24, 40, 6, 2, 2),
(3, 40, 80, 6, 2, 3),
(5, 80, 112, 6, 1, 3),
(5, 112, 192, 6, 2, 4),
(3, 192, 320, 6, 1, 1),
# k, c, e, s, n, r, a
(3, 16, 1, 1, 1, 4, swish),
(3, 24, 6, 2, 2, 4, swish),
(5, 40, 6, 2, 2, 4, swish),
(3, 80, 6, 2, 3, 4, swish),
(5, 112, 6, 1, 3, 4, swish),
(5, 192, 6, 2, 4, 4, swish),
(3, 320, 6, 1, 1, 4, swish),
]
# Data is organised as (r, (w, d))
# r: image resolution
Expand Down Expand Up @@ -46,7 +46,7 @@ function EfficientNet(config::Symbol; pretrain::Bool = false, inchannels::Intege
scalings = EFFICIENTNET_GLOBAL_CONFIGS[config][2]
layers = efficientnet(EFFICIENTNET_BLOCK_CONFIGS,
fill(mbconv_builder, length(EFFICIENTNET_BLOCK_CONFIGS));
scalings, inchannels, nclasses)
inplanes = 32, scalings, inchannels, nclasses)
if pretrain
loadpretrain!(layers, string("efficientnet-", config))
end
Expand Down
61 changes: 31 additions & 30 deletions src/convnets/efficientnets/efficientnetv2.jl
View file
Original file line number Diff line number Diff line change
@@ -1,36 +1,36 @@
# block configs for EfficientNetv2
# data organised as (k, i, o, e, s, n)
# data organised as (k, c, e, s, n, r, a)
const EFFNETV2_CONFIGS = Dict(:small => [
(3, 24, 24, 1, 1, 2),
(3, 24, 48, 4, 2, 4),
(3, 48, 64, 4, 2, 4),
(3, 64, 128, 4, 2, 6),
(3, 128, 160, 6, 1, 9),
(3, 160, 256, 6, 2, 15)],
(3, 24, 1, 1, 2, nothing, swish),
(3, 48, 4, 2, 4, nothing, swish),
(3, 64, 4, 2, 4, nothing, swish),
(3, 128, 4, 2, 6, 4, swish),
(3, 160, 6, 1, 9, 4, swish),
(3, 256, 6, 2, 15, 4, swish)],
:medium => [
(3, 24, 24, 1, 1, 3),
(3, 24, 48, 4, 2, 5),
(3, 48, 80, 4, 2, 5),
(3, 80, 160, 4, 2, 7),
(3, 160, 176, 6, 1, 14),
(3, 176, 304, 6, 2, 18),
(3, 304, 512, 6, 1, 5)],
(3, 24, 1, 1, 3, nothing, swish),
(3, 48, 4, 2, 5, nothing, swish),
(3, 80, 4, 2, 5, nothing, swish),
(3, 160, 4, 2, 7, 4, swish),
(3, 176, 6, 1, 14, 4, swish),
(3, 304, 6, 2, 18, 4, swish),
(3, 512, 6, 1, 5, 4, swish)],
:large => [
(3, 32, 32, 1, 1, 4),
(3, 32, 64, 4, 2, 7),
(3, 64, 96, 4, 2, 7),
(3, 96, 192, 4, 2, 10),
(3, 192, 224, 6, 1, 19),
(3, 224, 384, 6, 2, 25),
(3, 384, 640, 6, 1, 7)],
(3, 32, 1, 1, 4, nothing, swish),
(3, 64, 4, 2, 7, nothing, swish),
(3, 96, 4, 2, 7, nothing, swish),
(3, 192, 4, 2, 10, 4, swish),
(3, 224, 6, 1, 19, 4, swish),
(3, 384, 6, 2, 25, 4, swish),
(3, 640, 6, 1, 7, 4, swish)],
:xlarge => [
(3, 32, 32, 1, 1, 4),
(3, 32, 64, 4, 2, 8),
(3, 64, 96, 4, 2, 8),
(3, 96, 192, 4, 2, 16),
(3, 192, 224, 6, 1, 24),
(3, 384, 512, 6, 2, 32),
(3, 512, 768, 6, 1, 8)])
(3, 32, 1, 1, 4, nothing, swish),
(3, 64, 4, 2, 8, nothing, swish),
(3, 96, 4, 2, 8, nothing, swish),
(3, 192, 4, 2, 16, 4, swish),
(3, 384, 6, 1, 24, 4, swish),
(3, 512, 6, 2, 32, 4, swish),
(3, 768, 6, 1, 8, 4, swish)])

"""
EfficientNetv2(config::Symbol; pretrain::Bool = false, width_mult::Real = 1,
Expand Down Expand Up @@ -58,9 +58,10 @@ function EfficientNetv2(config::Symbol; pretrain::Bool = false,
layers = efficientnet(EFFNETV2_CONFIGS[config],
vcat(fill(fused_mbconv_builder, 3),
fill(mbconv_builder, length(EFFNETV2_CONFIGS[config]) - 3));
headplanes = 1280, inchannels, nclasses)
inplanes = EFFNETV2_CONFIGS[config][1][2], headplanes = 1280,
inchannels, nclasses)
if pretrain
loadpretrain!(layers, string("efficientnetv2"))
loadpretrain!(layers, string("efficientnetv2-", config))
end
return EfficientNetv2(layers)
end
Expand Down
43 changes: 20 additions & 23 deletions src/convnets/mobilenets/mobilenetv2.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,6 @@ Create a MobileNetv2 model.
+ `c`: The number of output feature maps
+ `n`: The number of times a block is repeated
+ `s`: The stride of the convolutional kernel
+ `a`: The activation function used in the bottleneck layer
- `width_mult`: Controls the number of output feature maps in each block
(with 1 being the default in the paper)
Expand All @@ -24,41 +23,39 @@ Create a MobileNetv2 model.
- `inchannels`: The number of input channels.
- `nclasses`: The number of output classes
"""
function mobilenetv2(configs::AbstractVector{<:Tuple}; width_mult::Real = 1,
max_width::Integer = 1280, divisor::Integer = 8, dropout_rate = 0.2,
function mobilenetv2(block_configs::AbstractVector{<:Tuple}; width_mult::Real = 1,
max_width::Integer = 1280, divisor::Integer = 8,
inplanes::Integer = 32, dropout_rate = 0.2,
inchannels::Integer = 3, nclasses::Integer = 1000)
# building first layer
inplanes = _round_channels(32 * width_mult, divisor)
inplanes = _round_channels(inplanes * width_mult, divisor)
layers = []
append!(layers,
conv_norm((3, 3), inchannels, inplanes; pad = 1, stride = 2))
# building inverted residual blocks
for (t, c, n, s, activation) in configs
outplanes = _round_channels(c * width_mult, divisor)
for i in 1:n
stride = i == 1 ? s : 1
push!(layers,
mbconv((3, 3), inplanes, round(Int, inplanes * t), outplanes,
activation; stride))
inplanes = outplanes
end
end
get_layers, block_repeats = mbconv_stack_builder(block_configs,
fill(mbconv_builder,
length(block_configs));
inplanes)
append!(layers, resnet_stages(get_layers, block_repeats, +))
# building last layers
outplanes = _round_channels(max_width * max(1, width_mult), divisor)
append!(layers, conv_norm((1, 1), inplanes, outplanes, relu6))
append!(layers,
conv_norm((1, 1), _round_channels(block_configs[end][2], 8),
outplanes, relu6))
return Chain(Chain(layers...), create_classifier(outplanes, nclasses; dropout_rate))
end

# Layer configurations for MobileNetv2
const MOBILENETV2_CONFIGS = [
# t, c, n, s, a
(1, 16, 1, 1, relu6),
(6, 24, 2, 2, relu6),
(6, 32, 3, 2, relu6),
(6, 64, 4, 2, relu6),
(6, 96, 3, 1, relu6),
(6, 160, 3, 2, relu6),
(6, 320, 1, 1, relu6),
# k, c, e, s, n, r, a
(3, 16, 1, 1, 1, nothing, relu6),
(3, 24, 6, 2, 2, nothing, relu6),
(3, 32, 6, 2, 3, nothing, relu6),
(3, 64, 6, 2, 4, nothing, relu6),
(3, 96, 6, 1, 3, nothing, relu6),
(3, 160, 6, 2, 3, nothing, relu6),
(3, 320, 6, 1, 1, nothing, relu6),
]

"""
Expand Down
2 changes: 1 addition & 1 deletion src/layers/mbconv.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -59,7 +59,7 @@ function fused_mbconv(kernel_size::Dims{2}, inplanes::Integer,
append!(layers, conv_norm((1, 1), explanes, outplanes, identity; norm_layer))
else
append!(layers,
conv_norm((1, 1), inplanes, outplanes, activation; norm_layer, stride))
conv_norm(kernel_size, inplanes, outplanes, activation; norm_layer, stride))
end
return Chain(layers...)
end

0 comments on commit 8e54d47

Please sign in to comment.