This commit is contained in:
@@ -0,0 +1,29 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import torch.nn as nn
|
||||||
|
|
||||||
|
from .resnet18_head import ResNet18Head
|
||||||
|
from .visual_syntax import VisualSyntaxTail
|
||||||
|
from core.data_models import VisualSyntaxData
|
||||||
|
|
||||||
|
|
||||||
|
class VisualCommunicationModel(nn.Module):
|
||||||
|
def __init__(self):
|
||||||
|
super().__init__()
|
||||||
|
# store other models
|
||||||
|
self.resnet_head = ResNet18Head()
|
||||||
|
self.visual_syntax_tail = VisualSyntaxTail()
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
# generate visual representation
|
||||||
|
vis_rep = self.resnet_head(x)
|
||||||
|
# prepare result map
|
||||||
|
results = {}
|
||||||
|
# predict visual syntax
|
||||||
|
visual_syntax_pred = self.visual_syntax_tail(vis_rep).cpu()
|
||||||
|
# visual_syntax_pred = float()
|
||||||
|
results['visual_syntax'] = VisualSyntaxData.from_list(
|
||||||
|
visual_syntax_pred,
|
||||||
|
)
|
||||||
|
|
||||||
|
return results
|
||||||
Reference in New Issue
Block a user