Update rfnet.py

Author: RONAK-AI647

examples/robot/lifelong_learning_bench/semantic-segmentation/testalgorithms/rfnet/RFNet/models/rfnet.py

@@ -1 +1,27 @@
-from logger import logger
+import torch.nn as nn
+from itertools import chain # 串联多个迭代对象
+from logger import logger
+from .util import _BNReluConv, upsample
+
+
+class RFNet(nn.Module):
+    def __init__(self, backbone, num_classes, use_bn=True):
+        super(RFNet, self).__init__()
+        self.backbone = backbone
+        self.num_classes = num_classes
+        logger.info(self.backbone.num_features)
+        self.logits = _BNReluConv(self.backbone.num_features, self.num_classes, batch_norm=use_bn)
+
+    def forward(self, rgb_inputs, depth_inputs = None):
+        x, additional = self.backbone(rgb_inputs, depth_inputs)
+        logits = self.logits.forward(x)
+        logits_upsample = upsample(logits, rgb_inputs.shape[2:])
+        #print(logits_upsample.size)
+        return logits_upsample
+
+
+    def random_init_params(self):
+        return chain(*([self.logits.parameters(), self.backbone.random_init_params()]))
+
+    def fine_tune_params(self):
+        return self.backbone.fine_tune_params()