def __init__(self, config: Config):
super().__init__()
self.input_frames = config['constants']['inputFrames']
self.metadata_dim = config['constants']['metadataDim']
self.embedding_dim = config['constants']['embeddingDim']
self.pre_squeeze = nn.Sequential(
nn.Conv2d(3 * self.input_frames,
12,
kernel_size=3,
stride=1,
padding=1), activation(inplace=True), nn.BatchNorm2d(12),
nn.Conv2d(12, 16, kernel_size=3, stride=1, padding=1, groups=2),
activation(inplace=True), nn.BatchNorm2d(16),
nn.Conv2d(16, 24, kernel_size=3, stride=2),
activation(inplace=True), nn.BatchNorm2d(24),
nn.Conv2d(24, 32, kernel_size=3, stride=2, groups=2),
activation(inplace=True), nn.BatchNorm2d(32), nn.Dropout2d(0.2))
self.squeeze_submodule = subnets.TestSqueeze(fire_func=Fire,
pool_func=pool_func)
self.post_squeeze = nn.Sequential(
nn.BatchNorm2d(128),
nn.Conv2d(128, 96, kernel_size=3, stride=2, padding=1),
activation(inplace=True),
nn.Conv2d(96, 64, kernel_size=3, stride=2, padding=1, groups=2),
activation(inplace=True),
nn.BatchNorm2d(64),
nn.Conv2d(64, 64, kernel_size=3, stride=2, padding=1),
activation(inplace=True),
nn.Conv2d(64, 64, kernel_size=3, stride=2, padding=1, groups=2),
)
self.meta_embedding = subnets.Metadata(8, self.metadata_dim)
self.post_metadata = nn.Sequential(
nn.Linear(128 + self.metadata_dim, 128),
activation(inplace=True),
nn.BatchNorm1d(128),
nn.Linear(128, 96),
activation(inplace=True),
nn.BatchNorm1d(96),
nn.Dropout(0.4),
nn.Linear(96, 96),
activation(inplace=True),
nn.BatchNorm1d(96),
nn.Linear(96, self.embedding_dim),
activation(inplace=True),
)
for mod in self.modules():
if hasattr(mod, 'weight') and hasattr(mod.weight, 'data'):
if isinstance(mod, nn.Conv2d):
init.kaiming_normal(mod.weight.data, a=-0.25)
elif len(mod.weight.data.size()) >= 2:
init.xavier_normal(mod.weight.data, gain=2.5)
else:
init.normal(mod.weight.data)
def __init__(self, config: Config):
"""Sets up layers"""
super().__init__()
self.input_frames = config['training']['inputFrames']
self.metadata_dim = config['constants']['metadataDim']
self.embedding_dim = config['constants']['embeddingDim']
self.pre_squeeze = nn.Sequential(
nn.Conv2d(6, 12, kernel_size=3, stride=1, padding=1),
activation(inplace=True), nn.BatchNorm2d(12),
nn.Conv2d(12, 16, kernel_size=3, stride=1, padding=1),
activation(inplace=True), nn.BatchNorm2d(16),
nn.Conv2d(16, 16, kernel_size=3, stride=2),
activation(inplace=True), nn.BatchNorm2d(16),
nn.AvgPool2d(kernel_size=3, stride=2), nn.Dropout2d(0.2))
self.squeeze_submodule = subnets.SqueezeSubmodule(fire_func=Fire,
pool_func=pool_func)
self.post_squeeze = nn.Sequential(
nn.BatchNorm2d(64),
nn.Conv2d(64, 32, kernel_size=3, stride=2, padding=1),
activation(inplace=True),
nn.Conv2d(32, 16, kernel_size=3, stride=2, padding=0, groups=1),
activation(inplace=True),
nn.BatchNorm2d(16),
nn.Conv2d(16, 16, kernel_size=2, stride=2, padding=1),
activation(inplace=True),
)
self.lstm_encoder = nn.ModuleList(
[nn.LSTM(32, 64, 1, batch_first=True)])
self.meta_embedding = subnets.Metadata(8, self.metadata_dim)
self.post_metadata = nn.Sequential(
nn.Linear(64 + self.metadata_dim, 64),
activation(inplace=True),
nn.BatchNorm1d(64),
nn.Linear(64, 64),
activation(inplace=True),
nn.BatchNorm1d(64),
nn.Dropout(0.4),
nn.Linear(64, 64),
activation(inplace=True),
nn.BatchNorm1d(64),
nn.Linear(64, self.embedding_dim),
activation(inplace=True),
)
for mod in self.modules():
if hasattr(mod, 'weight') and hasattr(mod.weight, 'data'):
if isinstance(mod, nn.Conv2d):
init.kaiming_normal(mod.weight.data, a=-0.25)
elif len(mod.weight.data.size()) >= 2:
init.xavier_normal(mod.weight.data, gain=2.5)
else:
init.normal(mod.weight.data)