def __init__(self, model_config: AttrDict, model_name: str):
super(EfficientNet, self).__init__()
assert model_config.INPUT_TYPE in ["rgb", "bgr"], "Input type not supported"
trunk_config = model_config.TRUNK.EFFICIENT_NETS
assert "model_version" in trunk_config, "Please specify EfficientNet version"
model_version = trunk_config["model_version"]
model_params = MODEL_PARAMS[model_version]
trunk_config["model_params"] = model_params
trunk_config.pop("model_version")
# we don't use the FC constructed with num_classes. This param is required
# to build the model in Classy Vision hence we pass the default value.
trunk_config["num_classes"] = 1000
logging.info(f"Building model: EfficientNet-{model_version}")
model = ClassyEfficientNet(**trunk_config)
self.drop_connect_rate = model.drop_connect_rate
self.num_blocks = len(model.blocks)
self.dropout = model.dropout
self.activation = Wrap(model.relu_fn) # using swish, not relu actually
# We map the layers of model into feature blocks to facilitate
# feature extraction at various layers of the model. The layers for which
# to extract features is controlled by out_feat_keys argument in the
# forward() call.
# - Stem
feature_blocks = [
["conv1", nn.Sequential(model.conv_stem, model.bn0, self.activation)]
]
# - Mobile Inverted Residual Bottleneck blocks
feature_blocks.extend(
[[f"block{i}", v] for i, v in enumerate(model.blocks.children())]
)
# - Conv Head + Pooling
feature_blocks.extend(
[
[
"conv_final",
nn.Sequential(model.conv_head, model.bn1, self.activation),
],
["avgpool", model.avg_pooling],
["flatten", Flatten(1)],
]
)
if model.dropout:
feature_blocks.append(["dropout", model.dropout])
# Consolidate into one indexable trunk
self._feature_blocks = nn.ModuleDict(feature_blocks)
self.all_feat_names = list(self._feature_blocks.keys())
def __init__(self, model_config: AttrDict, model_name: str):
super().__init__()
self.model_config = model_config
self.use_activation_checkpointing = (
model_config.ACTIVATION_CHECKPOINTING.USE_ACTIVATION_CHECKPOINTING
)
self.activation_checkpointing_splits = (
model_config.ACTIVATION_CHECKPOINTING.NUM_ACTIVATION_CHECKPOINTING_SPLITS
)
if self.use_activation_checkpointing:
logging.info(
f"Using Activation checkpointing. {self.activation_checkpointing_splits} chunks"
)
assert model_config.INPUT_TYPE in ["rgb", "bgr"], "Input type not supported"
trunk_config = model_config.TRUNK.REGNET
if "name" in trunk_config:
name = trunk_config["name"]
if name == "anynet":
model = build_model(trunk_config)
else:
logging.info(f"Building model: RegNet: {name}")
model = build_model({"name": name})
else:
logging.info("Building model: RegNet from yaml config")
model = ClassyRegNet.from_config(trunk_config)
# Now map the models to the structure we want to expose for SSL tasks
# The upstream RegNet model is made of :
# - `stem`
# - n x blocks in trunk_output, named `block1, block2, ..`
# We're only interested in the stem and successive blocks
# everything else is not picked up on purpose
feature_blocks: List[Tuple[str, nn.Module]] = []
# - get the stem
feature_blocks.append(("conv1", model.stem))
# - get all the feature blocks
for k, v in model.trunk_output.named_children():
assert k.startswith("block"), f"Unexpected layer name {k}"
block_index = len(feature_blocks) + 1
feature_blocks.append((f"res{block_index}", v))
# - finally, add avgpool and flatten.
feature_blocks.append(("avgpool", nn.AdaptiveAvgPool2d((1, 1))))
feature_blocks.append(("flatten", Flatten(1)))
self._feature_blocks = nn.ModuleDict(feature_blocks)
def __init__(self, model_config: AttrDict, model_name: str):
super().__init__()
conv1 = nn.Sequential(
nn.Conv2d(3, 64, kernel_size=11, stride=4, padding=2),
nn.BatchNorm2d(64),
nn.ReLU(inplace=True),
)
pool1 = nn.MaxPool2d(kernel_size=3, stride=2)
conv2 = nn.Sequential(
nn.Conv2d(64, 192, kernel_size=5, padding=2),
nn.BatchNorm2d(192),
nn.ReLU(inplace=True),
)
pool2 = nn.MaxPool2d(kernel_size=3, stride=2)
conv3 = nn.Sequential(
nn.Conv2d(192, 384, kernel_size=3, padding=1),
nn.BatchNorm2d(384),
nn.ReLU(inplace=True),
)
conv4 = nn.Sequential(
nn.Conv2d(384, 256, kernel_size=3, padding=1),
nn.BatchNorm2d(256),
nn.ReLU(inplace=True),
)
conv5 = nn.Sequential(
nn.Conv2d(256, 256, kernel_size=3, padding=1),
nn.BatchNorm2d(256),
nn.ReLU(inplace=True),
)
pool5 = nn.MaxPool2d(kernel_size=3, stride=2)
flatten = Flatten()
self._feature_blocks = nn.ModuleList(
[conv1, pool1, conv2, pool2, conv3, conv4, conv5, pool5, flatten]
)
self.all_feat_names = [
"conv1",
"pool1",
"conv2",
"pool2",
"conv3",
"conv4",
"conv5",
"pool5",
"flatten",
]
assert len(self.all_feat_names) == len(self._feature_blocks)
def __init__(
self,
model_config: AttrDict,
model_name: str,
):
"""
Args:
in_chans: Number of channels in the input to the U-Net model.
out_chans: Number of channels in the output to the U-Net model.
chans: Number of output channels of the first convolution layer.
num_pool_layers: Number of down-sampling and up-sampling layers.
drop_prob: Dropout probability.
"""
super().__init__()
self.model_config = model_config
trunk_config = model_config.TRUNK.TRUNK_PARAMS.UNET
in_chans = trunk_config.IN_CHANNELS
out_chans = trunk_config.OUT_CHANNELS
chans = trunk_config.get("CHANNELS", 32)
num_pool_layers = trunk_config.get("NUM_POOLS_LAYERS", 4)
drop_prob = trunk_config.get("DROP_PROBABILITY", 0.0)
self.use_checkpointing = (self.model_config.ACTIVATION_CHECKPOINTING.
USE_ACTIVATION_CHECKPOINTING)
self.num_checkpointing_splits = (
self.model_config.ACTIVATION_CHECKPOINTING.
NUM_ACTIVATION_CHECKPOINTING_SPLITS)
unetblock = UNetBlock(in_chans, out_chans, chans, num_pool_layers,
drop_prob)
## we mapped the layers of resnet model into feature blocks to facilitate
# feature extraction at various layers of the model. The layers for which
# to extract features is controlled by requested_feat_keys argument in the
# forward() call.
self._feature_blocks = nn.ModuleDict([
("unetblock", unetblock),
("flatten", Flatten(1)),
])
# give a name mapping to the layers so we can use a common terminology
# across models for feature evaluation purposes.
self.feat_eval_mapping = {
"res5": "unetblock",
"flatten": "flatten",
}
def __init__(self, model_config: AttrDict, model_name: str):
super().__init__()
self.model_config = model_config
trunk_config = model_config.TRUNK.TRUNK_PARAMS.NESTEDUNET
self.in_chans = trunk_config.IN_CHANNELS
self.out_chans = trunk_config.OUT_CHANNELS
self.deep_supervision = trunk_config.get("DEEP_SUPERVISION", False)
self.use_checkpointing = (self.model_config.ACTIVATION_CHECKPOINTING.
USE_ACTIVATION_CHECKPOINTING)
self.num_checkpointing_splits = (
self.model_config.ACTIVATION_CHECKPOINTING.
NUM_ACTIVATION_CHECKPOINTING_SPLITS)
self.nublock = NUBlock(self.in_chans, self.out_chans,
self.deep_supervision)
# give a name mapping to the layers so we can use a common terminology
# across models for feature evaluation purposes.
self.feat_eval_mapping = {
"nublock": "nublock",
}
# we mapped the layers of resnet model into feature blocks to facilitate
# feature extraction at various layers of the model. The layers for which
# to extract features is controlled by requested_feat_keys argument in the
# forward() call.
self._feature_blocks = nn.ModuleDict([
("nublock", self.nublock),
("flatten", Flatten(1)),
])
# give a name mapping to the layers so we can use a common terminology
# across models for feature evaluation purposes.
self.feat_eval_mapping = {
"nublock": "nublock",
"flatten": "flatten",
}
def __init__(self, model_config: AttrDict, model_name: str):
super().__init__()
conv1_relu = nn.Sequential(
nn.Conv2d(3, 96, kernel_size=11, stride=4, padding=0),
nn.ReLU(inplace=True))
lrn_pool1 = nn.Sequential(
nn.LocalResponseNorm(size=5, alpha=0.0001, beta=0.75),
nn.MaxPool2d(kernel_size=3, stride=2),
)
conv2_relu = nn.Sequential(
nn.Conv2d(96, 256, kernel_size=5, padding=2, groups=2),
nn.ReLU(inplace=True),
)
lrn_pool2 = nn.Sequential(
nn.LocalResponseNorm(size=5, alpha=0.0001, beta=0.75),
nn.MaxPool2d(kernel_size=3, stride=2),
)
conv3_relu = nn.Sequential(
nn.Conv2d(256, 384, kernel_size=3, padding=1),
nn.ReLU(inplace=True))
conv4_relu = nn.Sequential(
nn.Conv2d(384, 384, kernel_size=3, padding=1, groups=2),
nn.ReLU(inplace=True),
)
conv5_relu = nn.Sequential(
nn.Conv2d(384, 256, kernel_size=3, padding=1, groups=2),
nn.ReLU(inplace=True),
)
maxpool3 = nn.MaxPool2d(kernel_size=3, stride=2)
flatten = Flatten()
self._feature_blocks = nn.ModuleList([
conv1_relu,
lrn_pool1,
conv2_relu,
lrn_pool2,
conv3_relu,
conv4_relu,
conv5_relu,
maxpool3,
flatten,
])
self.all_feat_names = [
"conv1",
"pool1",
"conv2",
"pool2",
"conv3",
"conv4",
"conv5",
"pool5",
"flatten",
]
assert len(self.all_feat_names) == len(self._feature_blocks)
def __init__(self, model_config: AttrDict, model_name: str):
super(ResNeXt, self).__init__()
self.model_config = model_config
logging.info(
"ResNeXT trunk, supports activation checkpointing. {}".format(
"Activated" if self.model_config.ACTIVATION_CHECKPOINTING.
USE_ACTIVATION_CHECKPOINTING else "Deactivated"))
self.trunk_config = self.model_config.TRUNK.TRUNK_PARAMS.RESNETS
self.depth = SUPPORTED_DEPTHS(self.trunk_config.DEPTH)
self.width_multiplier = self.trunk_config.WIDTH_MULTIPLIER
self._norm_layer = _get_norm(self.trunk_config)
self.groups = self.trunk_config.GROUPS
self.zero_init_residual = self.trunk_config.ZERO_INIT_RESIDUAL
self.width_per_group = self.trunk_config.WIDTH_PER_GROUP
self.use_checkpointing = (self.model_config.ACTIVATION_CHECKPOINTING.
USE_ACTIVATION_CHECKPOINTING)
self.num_checkpointing_splits = (
self.model_config.ACTIVATION_CHECKPOINTING.
NUM_ACTIVATION_CHECKPOINTING_SPLITS)
(n1, n2, n3, n4) = BLOCK_CONFIG[self.depth]
logging.info(f"Building model: ResNeXt"
f"{self.depth}-{self.groups}x{self.width_per_group}d-"
f"w{self.width_multiplier}-{self._norm_layer.__name__}")
model = models.resnet.ResNet(
block=Bottleneck,
layers=(n1, n2, n3, n4),
zero_init_residual=self.zero_init_residual,
groups=self.groups,
width_per_group=self.width_per_group,
norm_layer=self._norm_layer,
)
model.inplanes = 64 * self.width_multiplier
dim_inner = 64 * self.width_multiplier
# some tasks like Colorization https://arxiv.org/abs/1603.08511 take input
# as L channel of an LAB image. In that case we change the input channel
# and re-construct the conv1
self.input_channels = INPUT_CHANNEL[self.model_config.INPUT_TYPE]
model_conv1 = nn.Conv2d(
self.input_channels,
model.inplanes,
kernel_size=7,
stride=2,
padding=3,
bias=False,
)
model_bn1 = self._norm_layer(model.inplanes)
model_relu1 = model.relu
model_maxpool = model.maxpool
model_avgpool = model.avgpool
model_layer1 = model._make_layer(Bottleneck, dim_inner, n1)
model_layer2 = model._make_layer(Bottleneck,
dim_inner * 2,
n2,
stride=2)
model_layer3 = model._make_layer(Bottleneck,
dim_inner * 4,
n3,
stride=2)
# For some models like Colorization https://arxiv.org/abs/1603.08511,
# due to the higher spatial resolution desired for pixel wise task, we
# support using a different stride. Currently, we know stride=1 and stride=2
# behavior so support only those.
safe_stride = SUPPORTED_L4_STRIDE(self.trunk_config.LAYER4_STRIDE)
model_layer4 = model._make_layer(Bottleneck,
dim_inner * 8,
n4,
stride=safe_stride)
# we mapped the layers of resnet model into feature blocks to facilitate
# feature extraction at various layers of the model. The layers for which
# to extract features is controlled by requested_feat_keys argument in the
# forward() call.
self._feature_blocks = nn.ModuleDict([
("conv1", model_conv1),
("bn1", model_bn1),
("conv1_relu", model_relu1),
("maxpool", model_maxpool),
("layer1", model_layer1),
("layer2", model_layer2),
("layer3", model_layer3),
("layer4", model_layer4),
("avgpool", model_avgpool),
("flatten", Flatten(1)),
])
# give a name mapping to the layers so we can use a common terminology
# across models for feature evaluation purposes.
self.feat_eval_mapping = {
"conv1": "conv1_relu",
"res1": "maxpool",
"res2": "layer1",
"res3": "layer2",
"res4": "layer3",
"res5": "layer4",
"res5avg": "avgpool",
"flatten": "flatten",
}
def create_regnet_feature_blocks(factory: RegnetBlocksFactory, model_config):
assert model_config.INPUT_TYPE in ["rgb",
"bgr"], "Input type not supported"
trunk_config = model_config.TRUNK.REGNET
if "name" in trunk_config:
assert (trunk_config["name"] == "anynet"
), "Please use AnyNetParams or specify RegNetParams dictionary"
if "name" in trunk_config and trunk_config["name"] == "anynet":
params = AnyNetParams(
depths=trunk_config["depths"],
widths=trunk_config["widths"],
group_widths=trunk_config["group_widths"],
bottleneck_multipliers=trunk_config["bottleneck_multipliers"],
strides=trunk_config["strides"],
stem_type=StemType[trunk_config.get("stem_type",
"simple_stem_in").upper()],
stem_width=trunk_config.get("stem_width", 32),
block_type=BlockType[trunk_config.get(
"block_type", "res_bottleneck_block").upper()],
activation=ActivationType[trunk_config.get("activation",
"relu").upper()],
use_se=trunk_config.get("use_se", True),
se_ratio=trunk_config.get("se_ratio", 0.25),
bn_epsilon=trunk_config.get("bn_epsilon", 1e-05),
bn_momentum=trunk_config.get("bn_momentum", 0.1),
)
else:
params = RegNetParams(
depth=trunk_config["depth"],
w_0=trunk_config["w_0"],
w_a=trunk_config["w_a"],
w_m=trunk_config["w_m"],
group_width=trunk_config["group_width"],
bottleneck_multiplier=trunk_config.get("bottleneck_multiplier",
1.0),
stem_type=StemType[trunk_config.get("stem_type",
"simple_stem_in").upper()],
stem_width=trunk_config.get("stem_width", 32),
block_type=BlockType[trunk_config.get(
"block_type", "res_bottleneck_block").upper()],
activation=ActivationType[trunk_config.get("activation",
"relu").upper()],
use_se=trunk_config.get("use_se", True),
se_ratio=trunk_config.get("se_ratio", 0.25),
bn_epsilon=trunk_config.get("bn_epsilon", 1e-05),
bn_momentum=trunk_config.get("bn_momentum", 0.1),
)
# Ad hoc stem
#
# Important: do NOT retain modules in self.stem or self.trunk_output. It may
# seem to be harmless, but it appears that autograd will result in computing
# grads in different order. Different ordering can cause deterministic OOM,
# even when the peak memory otherwise is only 24GB out of 32GB.
#
# When debugging this, it is not enough to just dump the total module
# params. You need to diff the module string representations.
stem = factory.create_stem(params)
# Instantiate all the AnyNet blocks in the trunk
current_width, trunk_depth, blocks = params.stem_width, 0, []
for i, (width_out, stride, depth, group_width,
bottleneck_multiplier) in enumerate(params.get_expanded_params()):
# Starting from 1
stage_index = i + 1
# Specify where the checkpoints are set in the stage
stage_checkpoints = []
if model_config.ACTIVATION_CHECKPOINTING.USE_ACTIVATION_CHECKPOINTING:
checkpoint_config = trunk_config.get("stage_checkpoints", [])
if len(checkpoint_config) > 0:
stage_checkpoints = checkpoint_config[i]
else:
stage_checkpoints.append(depth)
# Create the stage and add it to the trunk
new_stage = factory.create_any_stage(
width_in=current_width,
width_out=width_out,
stride=stride,
depth=depth,
group_width=group_width,
bottleneck_multiplier=bottleneck_multiplier,
params=params,
stage_index=stage_index,
checkpoints=stage_checkpoints,
)
blocks.append((f"block{stage_index}", new_stage))
trunk_depth += blocks[-1][1].stage_depth
current_width = width_out
trunk_output = nn.Sequential(OrderedDict(blocks))
################################################################################
# Now map the models to the structure we want to expose for SSL tasks
# The upstream RegNet model is made of :
# - `stem`
# - n x blocks in trunk_output, named `block1, block2, ..`
# We're only interested in the stem and successive blocks
# everything else is not picked up on purpose
feature_blocks: List[Tuple[str, nn.Module]] = [("conv1", stem)]
for k, v in trunk_output.named_children():
assert k.startswith("block"), f"Unexpected layer name {k}"
block_index = len(feature_blocks) + 1
feature_blocks.append((f"res{block_index}", v))
feature_blocks.append(("avgpool", nn.AdaptiveAvgPool2d((1, 1))))
feature_blocks.append(("flatten", Flatten(1)))
return nn.ModuleDict(feature_blocks), trunk_depth
def __init__(self, model_config: AttrDict, model_name: str):
super().__init__()
self.model_config = model_config
assert model_config.INPUT_TYPE in ["rgb",
"bgr"], "Input type not supported"
trunk_config = model_config.TRUNK.TRUNK_PARAMS.REGNET
assert "name" not in trunk_config, "Please specify the RegNet Params dictionary"
################################################################################
params = RegNetParams(
depth=trunk_config["depth"],
w_0=trunk_config["w_0"],
w_a=trunk_config["w_a"],
w_m=trunk_config["w_m"],
group_w=trunk_config["group_width"],
stem_type=trunk_config.get("stem_type", "simple_stem_in").upper(),
stem_width=trunk_config.get("stem_width", 32),
block_type=trunk_config.get("block_type",
"res_bottleneck_block").upper(),
use_se=trunk_config.get("use_se", True),
se_ratio=trunk_config.get("se_ratio", 0.25),
bn_epsilon=trunk_config.get("bn_epsilon", 1e-05),
bn_momentum=trunk_config.get("bn_momentum", 0.1),
)
# We need all workers (on all nodes) to have the same weights.
# Unlike DDP, FSDP does not sync weights using rank 0 on start.
# Therefore, we init stem and trunk_output below within the seed context.
#
# TODO (Min): we can make this seed coming from the config.
stem = None
trunk_output = None
with set_torch_seed(0):
# Ad hoc stem
#
# Important: do NOT retain modules in self.stem or self.trunk_output. It may
# seem to be harmless, but it appears that autograd will result in computing
# grads in different order. Different ordering can cause deterministic OOM,
# even when the peak memory otherwise is only 24GB out of 32GB.
#
# When debugging this, it is not enough to just dump the total module
# params. You need to diff the module string representations.
stem = {
StemType.RES_STEM_CIFAR: ResStemCifar,
StemType.RES_STEM_IN: ResStemIN,
StemType.SIMPLE_STEM_IN: SimpleStemIN,
}[params.stem_type](
3,
params.stem_width,
params.bn_epsilon,
params.bn_momentum,
params.relu_in_place,
)
init_weights(stem)
# Instantiate all the AnyNet blocks in the trunk
block_fun = {
BlockType.VANILLA_BLOCK: VanillaBlock,
BlockType.RES_BASIC_BLOCK: ResBasicBlock,
BlockType.RES_BOTTLENECK_BLOCK: ResBottleneckBlock,
}[params.block_type]
current_width = params.stem_width
self.trunk_depth = 0
blocks = []
for i, (width_out, stride, depth, bot_mul,
group_width) in enumerate(params.get_expanded_params()):
blocks.append((
f"block{i+1}",
AnyStage(
current_width,
width_out,
stride,
depth,
block_fun,
bot_mul,
group_width,
params,
stage_index=i + 1,
),
))
self.trunk_depth += blocks[-1][1].stage_depth
current_width = width_out
trunk_output = nn.Sequential(OrderedDict(blocks))
################################################################################
# Now map the models to the structure we want to expose for SSL tasks
# The upstream RegNet model is made of :
# - `stem`
# - n x blocks in trunk_output, named `block1, block2, ..`
# We're only interested in the stem and successive blocks
# everything else is not picked up on purpose
feature_blocks: List[Tuple[str, nn.Module]] = []
# - get the stem
feature_blocks.append(("conv1", stem))
# - get all the feature blocks
for k, v in trunk_output.named_children():
assert k.startswith("block"), f"Unexpected layer name {k}"
block_index = len(feature_blocks) + 1
feature_blocks.append((f"res{block_index}", v))
# - finally, add avgpool and flatten.
feature_blocks.append(("avgpool", nn.AdaptiveAvgPool2d((1, 1))))
feature_blocks.append(("flatten", Flatten(1)))
self._feature_blocks = nn.ModuleDict(feature_blocks)
def __init__(self, model_config: AttrDict, model_name: str):
super().__init__()
# first setup the sobel filter
grayscale = nn.Conv2d(3, 1, kernel_size=1, stride=1, padding=0)
grayscale.weight.data.fill_(1.0 / 3.0)
grayscale.bias.data.zero_()
sobel_filter = nn.Conv2d(1, 2, kernel_size=3, stride=1, padding=1)
sobel_filter.weight.data[0, 0].copy_(
torch.FloatTensor([[1, 0, -1], [2, 0, -2], [1, 0, -1]])
)
sobel_filter.weight.data[1, 0].copy_(
torch.FloatTensor([[1, 2, 1], [0, 0, 0], [-1, -2, -1]])
)
sobel_filter.bias.data.zero_()
self.sobel = nn.Sequential(grayscale, sobel_filter)
for p in self.sobel.parameters():
p.requires_grad = False
# Setup the features
conv1_bn_relu = nn.Sequential(
nn.Conv2d(2, 96, kernel_size=11, stride=4, padding=2),
nn.BatchNorm2d(96),
nn.ReLU(inplace=True),
)
pool1 = nn.MaxPool2d(kernel_size=3, stride=2)
conv2_bn_relu = nn.Sequential(
nn.Conv2d(96, 256, kernel_size=5, stride=1, padding=2),
nn.BatchNorm2d(256),
nn.ReLU(inplace=True),
)
pool2 = nn.MaxPool2d(kernel_size=3, stride=2)
conv3_bn_relu = nn.Sequential(
nn.Conv2d(256, 384, kernel_size=3, stride=1, padding=1),
nn.BatchNorm2d(384),
nn.ReLU(inplace=True),
)
conv4_bn_relu = nn.Sequential(
nn.Conv2d(384, 384, kernel_size=3, stride=1, padding=1),
nn.BatchNorm2d(384),
nn.ReLU(inplace=True),
)
conv5_bn_relu = nn.Sequential(
nn.Conv2d(384, 256, kernel_size=3, stride=1, padding=1),
nn.BatchNorm2d(256),
nn.ReLU(inplace=True),
)
pool3 = nn.MaxPool2d(kernel_size=3, stride=2)
flatten = Flatten()
self._feature_blocks = nn.ModuleList(
[
conv1_bn_relu,
pool1,
conv2_bn_relu,
pool2,
conv3_bn_relu,
conv4_bn_relu,
conv5_bn_relu,
pool3,
flatten,
]
)
self.all_feat_names = [
"conv1",
"pool1",
"conv2",
"pool2",
"conv3",
"conv4",
"conv5",
"pool5",
"flatten",
]
assert len(self.all_feat_names) == len(self._feature_blocks)
def __init__(self, model_config: AttrDict, model_name: str):
super().__init__()
conv1_bn_relu = nn.Sequential(
nn.Conv2d(1, 96, kernel_size=11, stride=4, padding=0),
nn.BatchNorm2d(96),
nn.ReLU(inplace=True),
)
pool1 = nn.MaxPool2d(kernel_size=3, stride=2)
conv2_bn_relu = nn.Sequential(
nn.Conv2d(96, 256, kernel_size=5, padding=2, groups=2),
nn.BatchNorm2d(256),
nn.ReLU(inplace=True),
)
pool2 = nn.MaxPool2d(kernel_size=3, stride=2)
conv3_bn_relu = nn.Sequential(
nn.Conv2d(256, 384, kernel_size=3, stride=1, padding=1),
nn.BatchNorm2d(384),
nn.ReLU(inplace=True),
)
conv4_bn_relu = nn.Sequential(
nn.Conv2d(384, 384, kernel_size=3, stride=1, padding=1, groups=2),
nn.BatchNorm2d(384),
nn.ReLU(inplace=True),
)
conv5_bn_relu = nn.Sequential(
nn.Conv2d(384, 256, kernel_size=3, stride=1, padding=1, groups=2),
nn.BatchNorm2d(256),
nn.ReLU(inplace=True),
)
pool3 = nn.MaxPool2d(kernel_size=3, stride=2)
flatten = Flatten()
self._feature_blocks = nn.ModuleList([
conv1_bn_relu,
pool1,
conv2_bn_relu,
pool2,
conv3_bn_relu,
conv4_bn_relu,
conv5_bn_relu,
pool3,
flatten,
])
self.all_feat_names = [
"conv1",
"pool1",
"conv2",
"pool2",
"conv3",
"conv4",
"conv5",
"pool5",
"flatten",
]
assert len(self.all_feat_names) == len(self._feature_blocks)
assert (model_config.INPUT_TYPE == "lab"
), "AlexNet Colorization model takes LAB image only"
def __init__(
self,
model_config: AttrDict,
model_name: str,
):
"""
Args:
in_chans: Number of channels in the input to the U-Net model.
out_chans: Number of channels in the output to the U-Net model.
chans: Number of output channels of the first convolution layer.
num_pool_layers: Number of down-sampling and up-sampling layers.
drop_prob: Dropout probability.
"""
super().__init__()
self.model_config = model_config
trunk_config = model_config.TRUNK.TRUNK_PARAMS.UNET
in_chans = trunk_config.IN_CHANNELS
out_chans = trunk_config.OUT_CHANNELS
chans = trunk_config.get("CHANNELS", 32)
num_pool_layers = trunk_config.get("NUM_POOLS_LAYERS", 4)
drop_prob = trunk_config.get("DROP_PROBABILITY", 0.0)
self.use_checkpointing = (self.model_config.ACTIVATION_CHECKPOINTING.
USE_ACTIVATION_CHECKPOINTING)
self.num_checkpointing_splits = (
self.model_config.ACTIVATION_CHECKPOINTING.
NUM_ACTIVATION_CHECKPOINTING_SPLITS)
down_sample_layers = nn.ModuleList(
[ConvBlock(in_chans, chans, drop_prob)])
ch = chans
for _ in range(num_pool_layers - 1):
down_sample_layers.append(ConvBlock(ch, ch * 2, drop_prob))
ch *= 2
conv = ConvBlock(ch, ch * 2, drop_prob)
up_conv = nn.ModuleList()
up_transpose_conv = nn.ModuleList()
for _ in range(num_pool_layers - 1):
up_transpose_conv.append(TransposeConvBlock(ch * 2, ch))
up_conv.append(ConvBlock(ch * 2, ch, drop_prob))
ch //= 2
up_transpose_conv.append(TransposeConvBlock(ch * 2, ch))
up_conv.append(
nn.Sequential(
ConvBlock(ch * 2, ch, drop_prob),
nn.Conv2d(ch, out_chans, kernel_size=1, stride=1),
))
# we mapped the layers of resnet model into feature blocks to facilitate
# feature extraction at various layers of the model. The layers for which
# to extract features is controlled by requested_feat_keys argument in the
# forward() call.
feature_blocks_mapping = []
for i in range(down_sample_layers.__len__()):
feature_blocks_mapping.append(
(f"downlayer{i+1}", down_sample_layers[i]))
feature_blocks_mapping.append(('conv', conv))
for i in range(up_transpose_conv.__len__()):
feature_blocks_mapping.append(
(f"uptranconvlayer{i+1}", up_transpose_conv[i]))
for i in range(up_conv.__len__()):
feature_blocks_mapping.append((f"upconvlayer{i+1}", up_conv[i]))
feature_blocks_mapping.append(("flatten", Flatten(1)))
self._feature_blocks = nn.ModuleDict(feature_blocks_mapping)
# give a name mapping to the layers so we can use a common terminology
# across models for feature evaluation purposes.
self.feat_eval_mapping = {
"conv1": "conv1_relu",
"res1": "maxpool",
"res2": "layer1",
"res3": "layer2",
"res4": "layer3",
"res5": "layer4",
"res5avg": "avgpool",
"flatten": "flatten",
}
def create_regnet_feature_blocks(factory: RegnetBlocksFactory, model_config):
assert model_config.INPUT_TYPE in ["rgb",
"bgr"], "Input type not supported"
trunk_config = model_config.TRUNK.TRUNK_PARAMS.REGNET
assert "name" not in trunk_config, "Please specify the RegNet Params dictionary"
params = RegNetParams(
depth=trunk_config["depth"],
w_0=trunk_config["w_0"],
w_a=trunk_config["w_a"],
w_m=trunk_config["w_m"],
group_width=trunk_config["group_width"],
stem_type=trunk_config.get("stem_type", "simple_stem_in").upper(),
stem_width=trunk_config.get("stem_width", 32),
block_type=trunk_config.get("block_type",
"res_bottleneck_block").upper(),
activation=trunk_config.get("activation_type", "relu").upper(),
use_se=trunk_config.get("use_se", True),
se_ratio=trunk_config.get("se_ratio", 0.25),
bn_epsilon=trunk_config.get("bn_epsilon", 1e-05),
bn_momentum=trunk_config.get("bn_momentum", 0.1),
)
# Ad hoc stem
#
# Important: do NOT retain modules in self.stem or self.trunk_output. It may
# seem to be harmless, but it appears that autograd will result in computing
# grads in different order. Different ordering can cause deterministic OOM,
# even when the peak memory otherwise is only 24GB out of 32GB.
#
# When debugging this, it is not enough to just dump the total module
# params. You need to diff the module string representations.
stem = factory.create_stem(params)
# Instantiate all the AnyNet blocks in the trunk
current_width = params.stem_width
trunk_depth = 0
blocks = []
for i, (width_out, stride, depth, bot_mul,
group_width) in enumerate(params.get_expanded_params()):
blocks.append((
f"block{i + 1}",
AnyStage(
factory=factory,
width_in=current_width,
width_out=width_out,
stride=stride,
depth=depth,
bot_mul=bot_mul,
group_width=group_width,
params=params,
stage_index=i + 1,
),
))
trunk_depth += blocks[-1][1].stage_depth
current_width = width_out
trunk_output = nn.Sequential(OrderedDict(blocks))
################################################################################
# Now map the models to the structure we want to expose for SSL tasks
# The upstream RegNet model is made of :
# - `stem`
# - n x blocks in trunk_output, named `block1, block2, ..`
# We're only interested in the stem and successive blocks
# everything else is not picked up on purpose
feature_blocks: List[Tuple[str, nn.Module]] = [("conv1", stem)]
for k, v in trunk_output.named_children():
assert k.startswith("block"), f"Unexpected layer name {k}"
block_index = len(feature_blocks) + 1
feature_blocks.append((f"res{block_index}", v))
feature_blocks.append(("avgpool", nn.AdaptiveAvgPool2d((1, 1))))
feature_blocks.append(("flatten", Flatten(1)))
return nn.ModuleDict(feature_blocks), trunk_depth