Spaces:
Running
on
Zero
Running
on
Zero
File size: 914 Bytes
e85fecb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 |
"""
Copyright (c) 2024 The D-FINE Authors. All Rights Reserved.
"""
import torch.nn as nn
from ...core import register
__all__ = [
"DFINE",
]
@register()
class DFINE(nn.Module):
__inject__ = [
"backbone",
"encoder",
"decoder",
]
def __init__(
self,
backbone: nn.Module,
encoder: nn.Module,
decoder: nn.Module,
):
super().__init__()
self.backbone = backbone
self.decoder = decoder
self.encoder = encoder
def forward(self, x, targets=None):
x = self.backbone(x)
x = self.encoder(x)
x = self.decoder(x, targets)
return x
def deploy(
self,
):
self.eval()
for m in self.modules():
if hasattr(m, "convert_to_deploy"):
m.convert_to_deploy()
return self
|