def __init__(self, cfg, num_modes=3, model_name='resnest50',
use_bn: bool = True,
hdim: int = 512,
pretrained=True,
in_channels: int = 0):
super(ResNeStMulti, self).__init__()
out_dim, num_preds, future_len = calc_out_channels(cfg, num_modes=num_modes)
self.in_channels = in_channels
self.out_dim = out_dim
self.num_preds = num_preds
self.future_len = future_len
self.num_modes = num_modes
# self.conv0 = nn.Conv2d(
# in_channels, 3, kernel_size=3, stride=1, padding=1, bias=True)
self.base_model = model_name_dict[model_name](pretrained=pretrained)
# --- Replace first conv block, instead of preparing conv0 ----
if isinstance(self.base_model.conv1, Sequential):
conv = self.base_model.conv1[0]
self.base_model.conv1[0] = nn.Conv2d(
in_channels, conv.out_channels, kernel_size=conv.kernel_size, stride=conv.stride,
padding=conv.padding, bias=True)
else:
conv = self.base_model.conv1
self.base_model.conv1 = nn.Conv2d(
in_channels, conv.out_channels, kernel_size=conv.kernel_size, stride=conv.stride,
padding=conv.padding, bias=True)
activation = F.leaky_relu
self.do_pooling = True
# if self.do_pooling:
# inch = self.base_model.last_linear.in_features
# else:
# inch = None
inch = None
lin1 = LinearBlock(inch, hdim, use_bn=use_bn, activation=activation, residual=False)
lin2 = LinearBlock(hdim, out_dim, use_bn=use_bn, activation=None, residual=False)
self.lin_layers = Sequential(lin1, lin2)
def __init__(
self,
cfg,
num_modes=3,
in_channels: int = 0,
hdim: int = 512,
use_bn: bool = False,
model_name: str = "smp_fpn",
encoder_name: str = "resnet18",
roi_kernel_size: float = 1.0,
):
super(SMPMultiAgentModel, self).__init__()
out_dim, num_preds, future_len = calc_out_channels(cfg, num_modes=num_modes)
if model_name == "smp_unet":
self.base_model = smp.Unet(encoder_name, in_channels=in_channels)
elif model_name == "smp_fpn":
self.base_model = smp.FPN(encoder_name, in_channels=in_channels)
else:
raise NotImplementedError(f"model_name {model_name} not supported in SMPMultiAgentModel")
# HACKING, skip conv2d to get
decoder_channels = self.base_model.segmentation_head[0].in_channels
print("decoder_channels", decoder_channels)
self.base_model.segmentation_head[0] = nn.Identity()
activation = F.leaky_relu
lin_head1 = LinearBlock(decoder_channels, hdim, use_bn=use_bn, activation=activation, residual=False)
lin_head2 = LinearBlock(hdim, out_dim, use_bn=use_bn, activation=None, residual=False)
self.lin_layers = Sequential(lin_head1, lin_head2)
# self.lin_head = nn.Linear(decoder_channels, out_channels)
self.in_channels = in_channels
self.model_name = model_name
self.num_preds = num_preds
self.out_dim = out_dim
self.future_len = future_len
self.num_modes = num_modes
self.roi_kernel_size = roi_kernel_size
def __init__(self,
cfg,
num_modes=3,
model_name="efficientnet-b0",
use_pretrained=True,
use_bn=True,
hdim: int = 512,
in_channels: int = 0):
super(EfficientNetMulti, self).__init__()
out_dim, num_preds, future_len = calc_out_channels(cfg,
num_modes=num_modes)
self.in_channels = in_channels
self.out_dim = out_dim
self.num_preds = num_preds
self.future_len = future_len
self.num_modes = num_modes
# self.conv0 = nn.Conv2d(
# in_channels, 3, kernel_size=3, stride=1, padding=1, bias=True)
self.base_model = EfficientNetWrapper(model_name=model_name,
use_pretrained=use_pretrained,
in_channels=in_channels)
activation = F.leaky_relu
inch = None
lin1 = LinearBlock(inch,
hdim,
use_bn=use_bn,
activation=activation,
residual=False)
lin2 = LinearBlock(hdim,
out_dim,
use_bn=use_bn,
activation=None,
residual=False)
self.lin_layers = Sequential(lin1, lin2)
def __init__(
self,
cfg,
in_channels: int,
num_modes=3,
backbone: str = "resnet18",
use_pretrained=True,
hdim: int = 4096,
feat_module_type: str = "none",
feat_channels: int = -1,
) -> None:
super().__init__()
out_dim, num_preds, future_len = calc_out_channels(cfg,
num_modes=num_modes)
self.in_channels = in_channels
self.num_modes = num_modes
self.hdim = hdim
self.out_dim = out_dim
self.num_preds = num_preds
self.future_len = future_len
self.backbone = timm.create_model(model_name=backbone,
pretrained=use_pretrained,
num_classes=0,
in_chans=in_channels)
self.avg_pool = torch.nn.AdaptiveAvgPool2d(1)
self.feature_dim = self.backbone.num_features
self.dense = nn.Sequential(nn.Linear(self.feature_dim, hdim),
nn.LeakyReLU(), nn.Linear(hdim, out_dim))
self.feat_module_type = feat_module_type
self.feat_module = FeatModule(
feat_module_type=feat_module_type,
channels=self.feature_dim,
feat_channels=feat_channels,
)
self.feat_channels = feat_channels
def __init__(
self,
cfg,
num_modes=3,
model_name='se_resnext101_32x4d',
use_bn: bool = True,
hdim: int = 512,
pretrained='imagenet',
in_channels: int = 0,
feat_module_type: str = "none",
feat_channels: int = -1,
):
super(PretrainedCNNMulti, self).__init__()
out_dim, num_preds, future_len = calc_out_channels(cfg,
num_modes=num_modes)
self.in_channels = in_channels
self.out_dim = out_dim
self.num_preds = num_preds
self.future_len = future_len
self.num_modes = num_modes
self.base_model = pretrainedmodels.__dict__[model_name](
pretrained=pretrained)
# --- Replace first conv ---
try:
if hasattr(self.base_model, "layer0") and isinstance(
self.base_model.layer0[0], nn.Conv2d):
print("Replace self.base_model.layer0[0]...")
# This works with SeResNeXt, but not tested with other network...
conv = self.base_model.layer0[0]
self.base_model.layer0[0] = nn.Conv2d(
in_channels,
conv.out_channels,
kernel_size=conv.kernel_size,
stride=conv.stride,
padding=conv.padding,
bias=True)
self.conv0 = None
# elif hasattr(self.base_model, "conv1") and isinstance(self.base_model.conv1, nn.Conv2d):
elif model_name in [
"resnet18", "resnet34", "resnet50", "resnet101",
"resnet152"
]:
# torchvision resnet is follows...
print("Replace base_model.conv1...")
self.base_model.conv1 = nn.Conv2d(
in_channels,
self.base_model.conv1.out_channels,
kernel_size=self.base_model.conv1.kernel_size,
stride=self.base_model.conv1.stride,
padding=self.base_model.conv1.padding,
bias=False,
)
self.conv0 = None
else:
raise ValueError("Cannot extract first conv layer")
except Exception as e:
# TODO: Better to replace `base_model`'s first conv block!
self.conv0 = nn.Conv2d(in_channels,
3,
kernel_size=3,
stride=1,
padding=1,
bias=True)
print(
f'[WARNING ]Cannot extract first conv layer for {model_name}, use conv0 to align channel size'
)
activation = F.leaky_relu
self.do_pooling = True
if self.do_pooling:
inch = self.base_model.last_linear.in_features
else:
inch = None
lin1 = LinearBlock(inch,
hdim,
use_bn=use_bn,
activation=activation,
residual=False)
lin2 = LinearBlock(hdim,
out_dim,
use_bn=use_bn,
activation=None,
residual=False)
self.lin_layers = Sequential(lin1, lin2)
self.feat_module_type = feat_module_type
self.feat_module = FeatModule(
feat_module_type=feat_module_type,
channels=inch,
feat_channels=feat_channels,
)
self.feat_channels = feat_channels