added model definitions
This commit is contained in:
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
Reference in New Issue
Block a user