重构Line_net2

lcnn/trainer.py

@@ -228,7 +228,7 @@ class Trainer(object):
             )
 
         if training:
-            self.model.train()
+            self.model.train1()
 
     def verify_freeze_params(model, freeze_config):
         """
@@ -253,7 +253,7 @@ class Trainer(object):
                                 print(f"  {param_name}: requires_grad = {param.requires_grad}")
 
     def train_epoch(self):
-        self.model.train()
+        self.model.train1()
 
         time = timer()
         for batch_idx, (image, meta, target, target_b) in enumerate(self.train_loader):

libs/vision_libs/__init__.py

@@ -3,7 +3,8 @@ import warnings
 from modulefinder import Module
 
 import torch
-from torchvision import _meta_registrations, datasets, io, models, ops, transforms, utils
+# from torchvision import _meta_registrations, datasets, io, models, ops, transforms, utils
+from torchvision import datasets, io, models, ops, transforms, utils
 
 from .extension import _HAS_OPS
 

models/base/base_detection_net.py

@@ -10,11 +10,9 @@ import torch
 from torch import nn, Tensor
 
 from libs.vision_libs.utils import _log_api_usage_once
-from models.base.base_model import BaseModel
-from models.line_detect.trainer import Trainer
 
 
-class BaseDetectionNet(BaseModel):
+class BaseDetectionNet(nn.Module):
     """
     Main class for Generalized R-CNN.
 
@@ -27,30 +25,6 @@ class BaseDetectionNet(BaseModel):
             the model
     """
 
-    def train(self, cfg):
-        pass
-
-    def get_loss(self, Loss, results, inputs, device):
-        pass
-
-    def get_optimizer(self, cfg_pipeline):
-        pass
-
-    def preprocess(self, cfg_pipeline):
-        pass
-
-    def transform(self, cfg_pipeline):
-        pass
-
-    def inference_begin(self, data):
-        pass
-
-    def inference_preprocess(self):
-        pass
-
-    def inference_end(self, inputs, results):
-        pass
-
     def __init__(self, backbone: nn.Module, rpn: nn.Module, roi_heads: nn.Module, transform: nn.Module) -> None:
         super().__init__()
         _log_api_usage_once(self)
@@ -69,13 +43,6 @@ class BaseDetectionNet(BaseModel):
 
         return detections
 
-    def train(self, cfg):
-        self.__trainer.train(self, "test")
-
-    def load_weight(self, pt_path):
-        state_dict = torch.load(pt_path)
-        self.load_state_dict(state_dict)
-
     def forward(self, images, targets=None):
         # type: (List[Tensor], Optional[List[Dict[str, Tensor]]]) -> Tuple[Dict[str, Tensor], List[Dict[str, Tensor]]]
         """
@@ -138,12 +105,12 @@ class BaseDetectionNet(BaseModel):
         proposals, proposal_losses = self.rpn(images, features, targets)
 
         detections, detector_losses = self.roi_heads(features, proposals, images.image_sizes, targets)
-        detections = self.transform.postprocess(detections, images.image_sizes,
-                                                original_image_sizes)  # type: ignore[operator]
+        detections = self.transform.postprocess(detections, images.image_sizes, original_image_sizes)  # type: ignore[operator]
 
         # ->multi task head
         # ->learner,->vectorize
 
+
         losses = {}
         losses.update(detector_losses)
         losses.update(proposal_losses)

models/base/base_model.py

@@ -14,93 +14,93 @@ class BaseModel(ABC, torch.nn.Module):
         self.trainer = None
 
     @abstractmethod
-    def train(self, cfg):
+    def train_by_cfg(self, cfg):
         return
 
-    @abstractmethod
-    def get_loss(self, Loss, results, inputs, device):
-        """Computes the loss given the network input and outputs.
-
-        Args:
-            Loss: A loss object.
-            results: This is the output of the model.
-            inputs: This is the input to the model.
-            device: The torch device to be used.
-
-        Returns:
-            Returns the loss value.
-        """
-        return
-
-    @abstractmethod
-    def get_optimizer(self, cfg_pipeline):
-        """Returns an optimizer object for the model.
-
-        Args:
-            cfg_pipeline: A Config object with the configuration of the pipeline.
-
-        Returns:
-            Returns a new optimizer object.
-        """
-        return
-
-    @abstractmethod
-    def preprocess(self, cfg_pipeline):
-        """Data preprocessing function.
-
-        This function is called before training to preprocess the data from a
-        dataset.
-
-        Args:
-            data: A sample from the dataset.
-            attr: The corresponding attributes.
-
-        Returns:
-            Returns the preprocessed data
-        """
-        return
-
-    @abstractmethod
-    def transform(self, cfg_pipeline):
-        """Transform function for the point cloud and features.
-
-        Args:
-            cfg_pipeline: config file for pipeline.
-        """
-        return
-
-    @abstractmethod
-    def inference_begin(self, data):
-        """Function called right before running inference.
-
-        Args:
-            data: A data from the dataset.
-        """
-        return
-
-    @abstractmethod
-    def inference_preprocess(self):
-        """This function prepares the inputs for the model.
-
-        Returns:
-            The inputs to be consumed by the call() function of the model.
-        """
-        return
-
-    @abstractmethod
-    def inference_end(self, inputs, results):
-        """This function is called after the inference.
-
-        This function can be implemented to apply post-processing on the
-        network outputs.
-
-        Args:
-            results: The model outputs as returned by the call() function.
-                Post-processing is applied on this object.
-
-        Returns:
-            Returns True if the inference is complete and otherwise False.
-            Returning False can be used to implement inference for large point
-            clouds which require multiple passes.
-        """
-        return
+    # @abstractmethod
+    # def get_loss(self, Loss, results, inputs, device):
+    #     """Computes the loss given the network input and outputs.
+    #
+    #     Args:
+    #         Loss: A loss object.
+    #         results: This is the output of the model.
+    #         inputs: This is the input to the model.
+    #         device: The torch device to be used.
+    #
+    #     Returns:
+    #         Returns the loss value.
+    #     """
+    #     return
+    #
+    # @abstractmethod
+    # def get_optimizer(self, cfg_pipeline):
+    #     """Returns an optimizer object for the model.
+    #
+    #     Args:
+    #         cfg_pipeline: A Config object with the configuration of the pipeline.
+    #
+    #     Returns:
+    #         Returns a new optimizer object.
+    #     """
+    #     return
+    #
+    # @abstractmethod
+    # def preprocess(self, cfg_pipeline):
+    #     """Data preprocessing function.
+    #
+    #     This function is called before training to preprocess the data from a
+    #     dataset.
+    #
+    #     Args:
+    #         data: A sample from the dataset.
+    #         attr: The corresponding attributes.
+    #
+    #     Returns:
+    #         Returns the preprocessed data
+    #     """
+    #     return
+    # #
+    # # @abstractmethod
+    # # def transform(self, cfg_pipeline):
+    # #     """Transform function for the point cloud and features.
+    # #
+    # #     Args:
+    # #         cfg_pipeline: config file for pipeline.
+    # #     """
+    # #     return
+    #
+    # @abstractmethod
+    # def inference_begin(self, data):
+    #     """Function called right before running inference.
+    #
+    #     Args:
+    #         data: A data from the dataset.
+    #     """
+    #     return
+    #
+    # @abstractmethod
+    # def inference_preprocess(self):
+    #     """This function prepares the inputs for the model.
+    #
+    #     Returns:
+    #         The inputs to be consumed by the call() function of the model.
+    #     """
+    #     return
+    #
+    # @abstractmethod
+    # def inference_end(self, inputs, results):
+    #     """This function is called after the inference.
+    #
+    #     This function can be implemented to apply post-processing on the
+    #     network outputs.
+    #
+    #     Args:
+    #         results: The model outputs as returned by the call() function.
+    #             Post-processing is applied on this object.
+    #
+    #     Returns:
+    #         Returns True if the inference is complete and otherwise False.
+    #         Returning False can be used to implement inference for large point
+    #         clouds which require multiple passes.
+    #     """
+    #     return

models/base/base_trainer.py

@@ -7,15 +7,15 @@ class BaseTrainer(ABC):
                  dataset=None,
                  device='cuda',
                  **kwargs):
-
+        #
         self.model = model
         self.dataset = dataset
         self.device=device
 
-    @abstractmethod
-    def train_cfg(self,model,cfg):
-        return
+    # @abstractmethod
+    # def train_cfg(self,model,cfg):
+    #     return
 
-    @abstractmethod
-    def train(self,model, **kwargs):
-        return
+    # @abstractmethod
+    # def train(self,model, **kwargs):
+    #     return

models/ins/trainer.py

@@ -14,7 +14,7 @@ from tools import utils, presets
 
 
 def train_one_epoch(model, optimizer, data_loader, device, epoch, print_freq, scaler=None):
-    model.train()
+    model.train1()
     metric_logger = utils.MetricLogger(delimiter="  ")
     metric_logger.add_meter("lr", utils.SmoothedValue(window_size=1, fmt="{value:.6f}"))
     header = f"Epoch: [{epoch}]"

models/ins_detect/trainer.py

@@ -14,7 +14,7 @@ from tools import utils, presets
 
 
 def train_one_epoch(model, optimizer, data_loader, device, epoch, print_freq, scaler=None):
-    model.train()
+    model.train1()
     metric_logger = utils.MetricLogger(delimiter="  ")
     metric_logger.add_meter("lr", utils.SmoothedValue(window_size=1, fmt="{value:.6f}"))
     header = f"Epoch: [{epoch}]"

models/keypoint/trainer.py

@@ -33,7 +33,7 @@ def log_losses_to_tensorboard(writer, result, step):
 
 
 def train_one_epoch(model, optimizer, data_loader, device, epoch, print_freq, writer, scaler=None):
-    model.train()
+    model.train1()
     metric_logger = utils.MetricLogger(delimiter="  ")
     metric_logger.add_meter("lr", utils.SmoothedValue(window_size=1, fmt="{value:.6f}"))
     header = f"Epoch: [{epoch}]"

models/line_detect/line_net.py

@@ -52,8 +52,8 @@ def _default_anchorgen():
 class LineNet(BaseDetectionNet):
     def __init__(self, cfg, **kwargs):
         cfg = read_yaml(cfg)
-        backbone = cfg['model']['backbone']
-        num_classes = cfg['model']['num_classes']
+        backbone = cfg['backbone']
+        num_classes = cfg['num_classes']
 
         if backbone == 'resnet50_fpn':
             is_trained = False
@@ -269,10 +269,11 @@ class LineNet(BaseDetectionNet):
     # self.roi_heads.line_head = line_head
     # self.roi_heads.line_predictor = line_predictor
 
-    def train(self, cfg):
+    def train_by_cfg(self, cfg):
         # cfg = read_yaml(cfg)
         self.trainer = Trainer()
-        self.trainer.train_cfg(model=self, cfg=cfg)
+        self.trainer.train_cfg(model=self,cfg=cfg)
+
 
 
 class TwoMLPHead(nn.Module):

models/line_detect/line_net.yaml

@@ -1,68 +1,57 @@
-
-
-
-
-model:
-  image:
-      mean: [109.730, 103.832, 98.681]
-      stddev: [22.275, 22.124, 23.229]
-
-  batch_size: 4
-  batch_size_eval: 2
-
-  # backbone multi-task parameters
-  head_size: [[2], [1], [2]]
-  loss_weight:
-    jmap: 8.0
-    lmap: 0.5
-    joff: 0.25
-    lpos: 1
-    lneg: 1
-    boxes: 1.0
-
-  # backbone parameters
-  backbone: resnet50_fpn
+image:
+  mean: [109.730, 103.832, 98.681]
+  stddev: [22.275, 22.124, 23.229]
+
+batch_size: 4
+batch_size_eval: 2
+
+# backbone multi-task parameters
+head_size: [[2], [1], [2]]
+loss_weight:
+jmap: 8.0
+lmap: 0.5
+joff: 0.25
+lpos: 1
+lneg: 1
+boxes: 1.0
+
+# backbone parameters
+backbone: resnet50_fpn
 #  backbone: unet
-  depth: 4
-  num_stacks: 1
-  num_blocks: 1
-  num_classes: 2
+depth: 4
+num_stacks: 1
+num_blocks: 1
+num_classes: 2
+
+# sampler parameters
+## static sampler
+n_stc_posl: 300
+n_stc_negl: 40
 
-  # sampler parameters
-  ## static sampler
-  n_stc_posl: 300
-  n_stc_negl: 40
+## dynamic sampler
+n_dyn_junc: 300
+n_dyn_posl: 300
+n_dyn_negl: 80
+n_dyn_othr: 600
 
-  ## dynamic sampler
-  n_dyn_junc: 300
-  n_dyn_posl: 300
-  n_dyn_negl: 80
-  n_dyn_othr: 600
+# LOIPool layer parameters
+n_pts0: 32
+n_pts1: 8
 
-  # LOIPool layer parameters
-  n_pts0: 32
-  n_pts1: 8
+# line verification network parameters
+dim_loi: 128
+dim_fc: 1024
 
-  # line verification network parameters
-  dim_loi: 128
-  dim_fc: 1024
+# maximum junction and line outputs
+n_out_junc: 250
+n_out_line: 2500
 
-  # maximum junction and line outputs
-  n_out_junc: 250
-  n_out_line: 2500
+# additional ablation study parameters
+use_cood: 0
+use_slop: 0
+use_conv: 0
 
-  # additional ablation study parameters
-  use_cood: 0
-  use_slop: 0
-  use_conv: 0
+# junction threashold for evaluation (See #5)
+eval_junc_thres: 0.008
 
-  # junction threashold for evaluation (See #5)
-  eval_junc_thres: 0.008
 
-optim:
-  name: Adam
-  lr: 4.0e-4
-  amsgrad: True
-  weight_decay: 1.0e-4
-  max_epoch: 1000
-  lr_decay_epoch: 10

models/line_detect/line_predictor.py

@@ -35,24 +35,24 @@ class LineRCNNPredictor(nn.Module):
         # self.cfg = read_yaml(cfg)
         # self.cfg = read_yaml(r'./config/wireframe.yaml')
         self.cfg = cfg
-        self.n_pts0 = self.cfg['model']['n_pts0']
-        self.n_pts1 = self.cfg['model']['n_pts1']
-        self.n_stc_posl = self.cfg['model']['n_stc_posl']
-        self.dim_loi = self.cfg['model']['dim_loi']
-        self.use_conv = self.cfg['model']['use_conv']
-        self.dim_fc = self.cfg['model']['dim_fc']
-        self.n_out_line = self.cfg['model']['n_out_line']
-        self.n_out_junc = self.cfg['model']['n_out_junc']
-        self.loss_weight = self.cfg['model']['loss_weight']
-        self.n_dyn_junc = self.cfg['model']['n_dyn_junc']
-        self.eval_junc_thres = self.cfg['model']['eval_junc_thres']
-        self.n_dyn_posl = self.cfg['model']['n_dyn_posl']
-        self.n_dyn_negl = self.cfg['model']['n_dyn_negl']
-        self.n_dyn_othr = self.cfg['model']['n_dyn_othr']
-        self.use_cood = self.cfg['model']['use_cood']
-        self.use_slop = self.cfg['model']['use_slop']
-        self.n_stc_negl = self.cfg['model']['n_stc_negl']
-        self.head_size = self.cfg['model']['head_size']
+        self.n_pts0 = self.cfg['n_pts0']
+        self.n_pts1 = self.cfg['n_pts1']
+        self.n_stc_posl = self.cfg['n_stc_posl']
+        self.dim_loi = self.cfg['dim_loi']
+        self.use_conv = self.cfg['use_conv']
+        self.dim_fc = self.cfg['dim_fc']
+        self.n_out_line = self.cfg['n_out_line']
+        self.n_out_junc = self.cfg['n_out_junc']
+        self.loss_weight = self.cfg['loss_weight']
+        self.n_dyn_junc = self.cfg['n_dyn_junc']
+        self.eval_junc_thres = self.cfg['eval_junc_thres']
+        self.n_dyn_posl = self.cfg['n_dyn_posl']
+        self.n_dyn_negl = self.cfg['n_dyn_negl']
+        self.n_dyn_othr = self.cfg['n_dyn_othr']
+        self.use_cood = self.cfg['use_cood']
+        self.use_slop = self.cfg['use_slop']
+        self.n_stc_negl = self.cfg['n_stc_negl']
+        self.head_size = self.cfg['head_size']
         self.num_class = sum(sum(self.head_size, []))
         self.head_off = np.cumsum([sum(h) for h in self.head_size])
 

models/line_detect/test_train.py

@@ -1,9 +1,12 @@
 import torch
 
 from models.line_detect.line_net import linenet_resnet50_fpn, LineNet
+from models.line_detect.trainer import Trainer
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 if __name__ == '__main__':
 
     model = LineNet('line_net.yaml')
-    model.train(cfg='./train.yaml')
+    # trainer = Trainer()
+    # trainer.train_cfg(model,cfg='./train.yaml')
+    model.train_by_cfg(cfg='train.yaml')

models/line_detect/train.yaml

@@ -1,67 +1,11 @@
 io:
   logdir: logs/
-  datadir: I:/datasets/wirenet_lm
+  datadir: D:\python\PycharmProjects\data
   resume_from:
   num_workers: 8
   tensorboard_port: 6000
   validation_interval: 300
 
-model:
-  image:
-      mean: [109.730, 103.832, 98.681]
-      stddev: [22.275, 22.124, 23.229]
-
-  batch_size: 4
-  batch_size_eval: 2
-
-  # backbone multi-task parameters
-  head_size: [[2], [1], [2]]
-  loss_weight:
-    jmap: 8.0
-    lmap: 0.5
-    joff: 0.25
-    lpos: 1
-    lneg: 1
-    boxes: 1.0
-
-  # backbone parameters
-  backbone: fasterrcnn_resnet50
-#  backbone: unet
-  depth: 4
-  num_stacks: 1
-  num_blocks: 1
-
-  # sampler parameters
-  ## static sampler
-  n_stc_posl: 300
-  n_stc_negl: 40
-
-  ## dynamic sampler
-  n_dyn_junc: 300
-  n_dyn_posl: 300
-  n_dyn_negl: 80
-  n_dyn_othr: 600
-
-  # LOIPool layer parameters
-  n_pts0: 32
-  n_pts1: 8
-
-  # line verification network parameters
-  dim_loi: 128
-  dim_fc: 1024
-
-  # maximum junction and line outputs
-  n_out_junc: 250
-  n_out_line: 2500
-
-  # additional ablation study parameters
-  use_cood: 0
-  use_slop: 0
-  use_conv: 0
-
-  # junction threashold for evaluation (See #5)
-  eval_junc_thres: 0.008
-
 optim:
   name: Adam
   lr: 4.0e-4

models/line_detect/trainer.py

@@ -1,6 +1,7 @@
 import torch
 from torch.utils.tensorboard import SummaryWriter
 
+from models.base.base_model import BaseModel
 from models.base.base_trainer import BaseTrainer
 from models.config.config_tool import read_yaml
 from models.line_detect.dataset_LD import WirePointDataset
@@ -21,13 +22,23 @@ def _loss(losses):
         total_loss += loss
 
     return total_loss
-
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+def move_to_device(data, device):
+    if isinstance(data, (list, tuple)):
+        return type(data)(move_to_device(item, device) for item in data)
+    elif isinstance(data, dict):
+        return {key: move_to_device(value, device) for key, value in data.items()}
+    elif isinstance(data, torch.Tensor):
+        return data.to(device)
+    else:
+        return data  # 对于非张量类型的数据不做任何改变
 
 class Trainer(BaseTrainer):
     def __init__(self, model=None,
                  dataset=None,
                  device='cuda',
                  **kwargs):
+
         super().__init__(model,dataset,device,**kwargs)
 
     def move_to_device(self, data, device):
@@ -54,15 +65,15 @@ class Trainer(BaseTrainer):
         except Exception as e:
             print(f"TensorBoard logging error: {e}")
 
-    def train_cfg(self, model, cfg):
+    def train_cfg(self, model:BaseModel, cfg):
         # cfg = r'./config/wireframe.yaml'
         cfg = read_yaml(cfg)
         print(f'cfg:{cfg}')
-        print(cfg['model']['n_dyn_negl'])
+        # print(cfg['n_dyn_negl'])
         self.train(model, **cfg)
 
-    def train(self, model, **cfg):
-        dataset_train = WirePointDataset(dataset_path=cfg['io']['datadir'], dataset_type='train')
+    def train(self, model, **kwargs):
+        dataset_train = WirePointDataset(dataset_path=kwargs['io']['datadir'], dataset_type='train')
         train_sampler = torch.utils.data.RandomSampler(dataset_train)
         # test_sampler = torch.utils.data.SequentialSampler(dataset_test)
         train_batch_sampler = torch.utils.data.BatchSampler(train_sampler, batch_size=2, drop_last=True)
@@ -71,7 +82,7 @@ class Trainer(BaseTrainer):
             dataset_train, batch_sampler=train_batch_sampler, num_workers=8, collate_fn=train_collate_fn
         )
 
-        dataset_val = WirePointDataset(dataset_path=cfg['io']['datadir'], dataset_type='val')
+        dataset_val = WirePointDataset(dataset_path=kwargs['io']['datadir'], dataset_type='val')
         val_sampler = torch.utils.data.RandomSampler(dataset_val)
         # test_sampler = torch.utils.data.SequentialSampler(dataset_test)
         val_batch_sampler = torch.utils.data.BatchSampler(val_sampler, batch_size=2, drop_last=True)
@@ -82,15 +93,19 @@ class Trainer(BaseTrainer):
 
         # model = linenet_resnet50_fpn().to(self.device)
 
-        optimizer = torch.optim.Adam(model.parameters(), lr=cfg['optim']['lr'])
-        writer = SummaryWriter(cfg['io']['logdir'])
+        optimizer = torch.optim.Adam(model.parameters(), lr=kwargs['optim']['lr'])
+        writer = SummaryWriter(kwargs['io']['logdir'])
 
-        for epoch in range(cfg['optim']['max_epoch']):
+        for epoch in range(kwargs['optim']['max_epoch']):
             print(f"epoch:{epoch}")
             model.train()
 
             for imgs, targets in data_loader_train:
-                losses = model(self.move_to_device(imgs, self.device), self.move_to_device(targets, self.device))
+                imgs = move_to_device(imgs, device)
+                targets=move_to_device(targets,device)
+                print(f'imgs:{len(imgs)}')
+                print(f'targets:{len(targets)}')
+                losses = model(imgs, targets)
                 # print(losses)
                 loss = _loss(losses)
                 optimizer.zero_grad()

models/wirenet2/trainer.py

@@ -13,7 +13,7 @@ from models.ins.maskrcnn_dataset import MaskRCNNDataset
 from models.keypoint.keypoint_dataset import KeypointDataset
 from tools import utils, presets
 def train_one_epoch(model, optimizer, data_loader, device, epoch, print_freq, scaler=None):
-    model.train()
+    model.train1()
     metric_logger = utils.MetricLogger(delimiter="  ")
     metric_logger.add_meter("lr", utils.SmoothedValue(window_size=1, fmt="{value:.6f}"))
     header = f"Epoch: [{epoch}]"

tools/engine.py

@@ -10,7 +10,7 @@ from tools.coco_utils import get_coco_api_from_dataset
 
 
 def train_one_epoch(model, optimizer, data_loader, device, epoch, print_freq, scaler=None):
-    model.train()
+    model.train1()
     metric_logger = utils.MetricLogger(delimiter="  ")
     metric_logger.add_meter("lr", utils.SmoothedValue(window_size=1, fmt="{value:.6f}"))
     header = f"Epoch: [{epoch}]"