added model definitions

This commit is contained in:
Brian Bjarke Jensen
2024-03-12 21:50:47 +01:00
parent ec48879327
commit f139206758
14 changed files with 114 additions and 0 deletions
View File
+8
View File
@@ -0,0 +1,8 @@
from __future__ import annotations
from .fully_connected import FullyConnectedModel
class AngleTail(FullyConnectedModel):
def __init__(self):
super().__init__(num_out_features=3)
+8
View File
@@ -0,0 +1,8 @@
from __future__ import annotations
from .fully_connected import FullyConnectedModel
class ContactTail(FullyConnectedModel):
def __init__(self):
super().__init__(num_out_features=2)
+8
View File
@@ -0,0 +1,8 @@
from __future__ import annotations
from .fully_connected import FullyConnectedModel
class DistanceTail(FullyConnectedModel):
def __init__(self):
super().__init__(num_out_features=3)
+8
View File
@@ -0,0 +1,8 @@
from __future__ import annotations
from .fully_connected import FullyConnectedModel
class FramingTail(FullyConnectedModel):
def __init__(self):
super().__init__(num_out_features=4)
+26
View File
@@ -0,0 +1,26 @@
from __future__ import annotations
import torch.nn as nn
class FullyConnectedModel(nn.Module):
def __init__(self, num_out_features: int):
super().__init__()
# define layers
self.fc1 = nn.Linear(in_features=16*16*512, out_features=512)
self.af1 = nn.ReLU()
self.fc2 = nn.Linear(in_features=512, out_features=128)
self.af2 = nn.ReLU()
self.fc3 = nn.Linear(in_features=128, out_features=32)
self.af3 = nn.ReLU()
self.fc4 = nn.Linear(in_features=32, out_features=num_out_features)
def forward(self, x):
x = self.fc1(x)
x = self.af1(x)
x = self.fc2(x)
x = self.af2(x)
x = self.fc3(x)
x = self.af3(x)
x = self.fc4(x)
return x
+8
View File
@@ -0,0 +1,8 @@
from __future__ import annotations
from .fully_connected import FullyConnectedModel
class InformationValueTail(FullyConnectedModel):
def __init__(self):
super().__init__(num_out_features=3)
+8
View File
@@ -0,0 +1,8 @@
from __future__ import annotations
from .fully_connected import FullyConnectedModel
class ModalityColorTail(FullyConnectedModel):
def __init__(self):
super().__init__(num_out_features=3)
+8
View File
@@ -0,0 +1,8 @@
from __future__ import annotations
from .fully_connected import FullyConnectedModel
class ModalityDepthTail(FullyConnectedModel):
def __init__(self):
super().__init__(num_out_features=3)
+8
View File
@@ -0,0 +1,8 @@
from __future__ import annotations
from .fully_connected import FullyConnectedModel
class ModalityLightingTail(FullyConnectedModel):
def __init__(self):
super().__init__(num_out_features=3)
+8
View File
@@ -0,0 +1,8 @@
from __future__ import annotations
from .fully_connected import FullyConnectedModel
class PointOfViewTail(FullyConnectedModel):
def __init__(self):
super().__init__(num_out_features=2)
+37
View File
@@ -0,0 +1,37 @@
from __future__ import annotations
import torch.nn as nn
import torchvision
class ResNet18Head(nn.Module):
def __init__(self):
super().__init__()
# copy out parts from ResNet18 with weights
resnet18 = torchvision.models.resnet18(
weights=torchvision.models.ResNet18_Weights.IMAGENET1K_V1,
)
# save relevant layers
self.conv1 = resnet18.conv1
self.bn1 = resnet18.bn1
self.relu = resnet18.relu
self.maxpool = resnet18.maxpool
self.layer1 = resnet18.layer1
self.layer2 = resnet18.layer2
self.layer3 = resnet18.layer3
self.layer4 = resnet18.layer4
self.avgpool = resnet18.avgpool
self.flat = nn.Flatten() # size 512
def forward(self, x):
x = self.conv1(x)
x = self.bn1(x)
x = self.relu(x)
x = self.maxpool(x)
x = self.layer1(x)
x = self.layer2(x)
x = self.layer3(x)
x = self.layer4(x)
x = self.avgpool(x)
x = self.flat(x)
return x
+8
View File
@@ -0,0 +1,8 @@
from __future__ import annotations
from .fully_connected import FullyConnectedModel
class SalienceTail(FullyConnectedModel):
def __init__(self):
super().__init__(num_out_features=5)
+8
View File
@@ -0,0 +1,8 @@
from __future__ import annotations
from .fully_connected import FullyConnectedModel
class VisualSyntaxTail(FullyConnectedModel):
def __init__(self):
super().__init__(num_out_features=18)