trying to get basic model working

2023-01-06 22:49:57 +00:00
parent cb104bd4d3
commit 93357f81cb
5 changed files with 143 additions and 17 deletions
@@ -3,15 +3,24 @@ import sys

 import torch
 import torch.nn as nn
+from torch.nn import functional as F
+
+sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
+from models.convnext import ConvNeXt


 class DepthEstimation(nn.Module):
    def __init__(self, device='cuda'):
        super().__init__()

-        self.layer_1 = nn.Conv2d(in_channels=3, out_channels=1, kernel_size=3, device=device, padding='same')  
+        self.convnext = ConvNeXt()
+        self.flatten = nn.Flatten()
+        self.decoder = ConvNeXt(reversed=True)
+

    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        x = self.layer_1(x)
+        x = self.convnext(x)
+        # x = self.flatten(x)
+        x = self.decoder(x)
        
        return x