kalpkanungo/SceneGraphNet

1import torch.nn as nn2 3class RelationshipNet(nn.Module):4    def __init__(self, num_classes):5        super().__init__()6 7        self.model = nn.Sequential(8            nn.Conv2d(3, 32, 3, padding=1),9            nn.ReLU(),10            nn.MaxPool2d(2),11 12            nn.Conv2d(32, 64, 3, padding=1),13            nn.ReLU(),14            nn.MaxPool2d(2),15 16            nn.Conv2d(64, 128, 3, padding=1),17            nn.ReLU(),18            nn.MaxPool2d(2),19 20            nn.Flatten(),21            nn.Linear(128 * 16 * 16, 256),22            nn.ReLU(),23            nn.Linear(256, num_classes)24        )25 26    def forward(self, x):27        return self.model(x)