添加根据val_loss保存权重功能

models/base/base_detection_net.py

@@ -41,10 +41,13 @@ class BaseDetectionNet(BaseModel):
     def eager_outputs(self, losses, detections):
         # type: (Dict[str, Tensor], List[Dict[str, Tensor]]) -> Union[Dict[str, Tensor], List[Dict[str, Tensor]]]
         if self.training:
-            return losses
-
+            return  losses
         return detections
 
+
+
+        # return detections,losses
+
     def forward(self, images, targets=None):
         # type: (List[Tensor], Optional[List[Dict[str, Tensor]]]) -> Tuple[Dict[str, Tensor], List[Dict[str, Tensor]]]
         """
@@ -59,6 +62,14 @@ class BaseDetectionNet(BaseModel):
                 like `scores`, `labels` and `mask` (for Mask R-CNN models).
 
         """
+        if targets is not None:
+            self.training = True
+            # print(f'targets is not None')
+
+        else:
+            self.training = False
+            # print(f'targets is None')
+
         if self.training:
             if targets is None:
                 torch._assert(False, "targets should not be none when in training mode")

models/line_detect/111.py

@@ -182,7 +182,8 @@ class Trainer(BaseTrainer):
         )
 
         last_model_path = os.path.join(wts_path, 'last.pth')
-        best_model_path = os.path.join(wts_path, 'best.pth')
+        best_train_model_path = os.path.join(wts_path, 'best_train.pth')
+        best_val_model_path = os.path.join(wts_path, 'best_val.pth')
         global_step = 0
 
         for epoch in range(kwargs['optim']['max_epoch']):
@@ -201,22 +202,28 @@ class Trainer(BaseTrainer):
                 self.writer_loss(writer, losses, global_step)
                 global_step += 1
 
-            avg_train_loss = total_train_loss / len(data_loader_train)
-            if epoch == 0:
-                best_loss = avg_train_loss
-            writer.add_scalar('loss/train', avg_train_loss, epoch)
 
-            if os.path.exists(f'{wts_path}/last.pt'):
-                os.remove(f'{wts_path}/last.pt')
-            save_last_model(model, last_model_path, epoch, optimizer)
-            best_loss = save_best_model(model, best_model_path, epoch, avg_train_loss, best_loss, optimizer)
 
             model.eval()
+            print(f'model.eval!!')
+            # ========== Validation ==========
+            total_val_loss = 0.0
             with torch.no_grad():
                 for batch_idx, (imgs, targets) in enumerate(data_loader_val):
                     t_start = time.time()
                     print(f'start to predict:{t_start}')
-                    pred = model(self.move_to_device(imgs, self.device))
+
+                    imgs = move_to_device(imgs, device)
+                    targets = move_to_device(targets, device)
+                    print(f'targets:{targets}')
+
+                    losses = model(imgs, targets)
+                    print(f'val losses:{losses}')
+                    loss = _loss(losses)
+                    total_val_loss += loss.item()
+
+                    pred= model(self.move_to_device(imgs, self.device))
+
                     # print(f'pred:{pred}')
                     t_end = time.time()
                     print(f'predict used:{t_end - t_start}')
@@ -224,6 +231,27 @@ class Trainer(BaseTrainer):
                         show_line(imgs[0], pred, epoch, writer)
                     break
 
+            avg_val_loss = total_val_loss / len(data_loader_val)
+            # print(f'avg_val_loss:{avg_val_loss}')
+
+            avg_train_loss = total_train_loss / len(data_loader_train)
+            print(f'avg_train_loss:{avg_train_loss}')
+            if epoch == 0:
+                best_train_loss = avg_train_loss
+                best_val_loss = avg_val_loss
+            writer.add_scalar('loss/train', avg_train_loss, epoch)
+
+            if os.path.exists(f'{wts_path}/last.pt'):
+                os.remove(f'{wts_path}/last.pt')
+            save_last_model(model, last_model_path, epoch, optimizer)
+            best_train_loss = save_best_model(model, best_train_model_path, epoch, avg_train_loss, best_train_loss,
+                                              optimizer)
+
+            best_val_loss = save_best_model(model, best_val_model_path, epoch, avg_val_loss, best_val_loss,
+                                              optimizer)
+            writer.add_scalar('loss/val', avg_val_loss, epoch)
+            print(f"Epoch {epoch} - Val Loss: {avg_val_loss:.4f}")
+
 
 import torch
 

models/line_detect/predict2.py

@@ -318,9 +318,9 @@ if __name__ == '__main__':
     model=get_line_net_convnext_fpn(num_classes=2).to(device)
     # pt_path = r"C:\Users\m2337\Downloads\best_lmap代替x，训练24轮结果.pth"
     # pt_path = r"C:\Users\m2337\Downloads\best_lmap代替x，训练75轮.pth"
-    pt_path = r"\\192.168.50.222\share\rlq\weights\convnext250514.pth"
+    pt_path = r"\\192.168.50.222\share\rlq\weights\convnext25051401.pth"
     # pt_path = r"C:\Users\m2337\Downloads\best_e20.pth"
-    img_path = r"color_img.jpg"
+    img_path = r"\\192.168.50.222\share\zyh\513\a_dataset\images\val\2025-05-13-08-56-03_LaserData_ID019504_color_scale.jpg"
     predict(pt_path, model, img_path)
     t_end = time.time()
     print(f'predict used:{t_end - t_start}')

models/line_detect/roi_heads.py

@@ -1002,9 +1002,11 @@ class RoIHeads(nn.Module):
         """
         if targets is not None:
             self.training = True
+            # print(f'targets is not None')
 
         else:
             self.training = False
+            # print(f'targets is None')
 
         if targets is not None:
             for t in targets:
@@ -1025,6 +1027,7 @@ class RoIHeads(nn.Module):
             regression_targets = None
             matched_idxs = None
 
+
         box_features = self.box_roi_pool(features, proposals, image_shapes)
         box_features = self.box_head(box_features)
         class_logits, box_regression = self.box_predictor(box_features)
@@ -1039,6 +1042,7 @@ class RoIHeads(nn.Module):
             loss_classifier, loss_box_reg = fastrcnn_loss(class_logits, box_regression, labels, regression_targets)
             losses = {"loss_classifier": loss_classifier, "loss_box_reg": loss_box_reg}
         else:
+
             boxes, scores, labels = self.postprocess_detections(class_logits, box_regression, proposals, image_shapes)
             num_images = len(boxes)
             for i in range(num_images):
@@ -1074,13 +1078,18 @@ class RoIHeads(nn.Module):
             if self.training:
                 rcnn_loss_wirepoint = wirepoint_head_line_loss(targets, line_features, x, y, idx, loss_weight)
                 loss_wirepoint = {"loss_wirepoint": rcnn_loss_wirepoint}
+                # print(f'loss_wirepoint:{loss_wirepoint}')
 
             else:
+                # rcnn_loss_wirepoint = wirepoint_head_line_loss(targets, line_features, x, y, idx, loss_weight)
 
+                print(f'model inference!!!')
                 pred = wirepoint_inference(x, idx, jcs, n_batch, ps, n_out_line, n_out_junc)
                 result.append(line_features)
                 result.append(pred)
+
                 loss_wirepoint = {}
+
             losses.update(loss_wirepoint)
         else:
             pass
@@ -1176,4 +1185,5 @@ class RoIHeads(nn.Module):
                     r["keypoints_scores"] = kps
             losses.update(loss_keypoint)
 
+        # print(f'roi losses:{losses}')
         return result, losses

utils/log_util.py

@@ -86,6 +86,7 @@ def save_best_model(model, save_path, epoch, current_loss, best_loss, optimizer=
 #     # Log the image to TensorBoard or other logger
 #     writer.add_image('validate', image_from_plot, epoch, dataformats='HWC')
 def show_line(img, pred, epoch, writer):
+    img=img.cpu().detach()
     im = img.permute(1, 2, 0)
     writer.add_image("z-ori", im, epoch, dataformats="HWC")