Training works

cli/fcos/__main__.py

@@ -23,6 +23,8 @@
 train_parser.add_argument('--optimizer_name', type=str, default="Adam",  required=False)
 train_parser.add_argument('--save_file', type=pathlib.Path, required=False)
 
+predict_parser = sub_parsers.add_parser("predict")
+
 args = parser.parse_args()
 
 if args.command == "train":
@@ -32,3 +34,5 @@
         train(**vars(args))
     except ImportError:
         print("Train Module not included.")
+elif args.command == "predict":
+    pass

core/fcos/core/mAP/functions.py

@@ -124,7 +124,7 @@ def calculate_mAP_gt(box, gt_box,threshold=0.5):
     iou = compute_iou(gt_box[1:5], box[2:6])
     if iou >= threshold:
         TP = True
-    return (box[:2], TP)
+    return (*box[:2], TP)
 
 def return_mAP(model, dataset, classes):
     model.cuda()
@@ -136,26 +136,27 @@ def return_mAP(model, dataset, classes):
     # traverse dataset
     for image, tags in dataset:
         # get an image
-        cuda_image = image.cuda()
+        cuda_image = image.unsqueeze(0).cuda()
         # predict
         row = cuda_image.shape[2]
         col = cuda_image.shape[3]
         confs, locs, centers = model(cuda_image)
-        boxes = fcos_to_boxes(confs, locs, centers, row, col)
+        boxes = fcos_to_boxes(classes, confs, locs, centers, row, col)
         for gt_box in tags:
-            gt_count_all[gt_box[0]] += 1
+            box_class = classes[int(gt_box[0].item())]
+            gt_count_all[box_class] += 1
             for box in boxes:
                 if(box[0] == gt_box[0]):
-                    mAP_all[box[0]].append(calculate_mAP_gt(box, gt_box))
+                    mAP_all[box_class].append(calculate_mAP_gt(box, gt_box))
     mAP = 0
-    # p = 0
-    # r = 0
-    for c in mAP_all:
+    mp = 0
+    mr = 0
+    for c in classes:
         p, r, ap = compute_mAP(mAP_all[c], gt_count_all[c])
         mAP += ap * gt_count_all[c]
-        # p += p * gt_count_all[c]
-        # r += r * gt_count_all[c]
+        mp += p * gt_count_all[c]
+        mr += r * gt_count_all[c]
     mAP /= sum(gt_count_all.values())
-    # p /= sum(gt_count_all.values())
-    # r /= sum(gt_count_all.values())
-    return mAP
+    mp /= sum(gt_count_all.values())
+    mr /= sum(gt_count_all.values())
+    return mAP, mp, mr

core/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "fcos-core"
-version = "0.1.4"
+version = "0.1.7"
 description = ""
 authors = ["DTheLegend <[email protected]>"]
 readme = "README.md"

train/fcos/train/train.py

@@ -6,6 +6,8 @@
 import os
 from fcos.core.mAP import return_mAP
 import torch
+from tqdm import tqdm
+from math import ceil
 
 # torch.manual_seed(1)	#reproducible
 def train(weights, classfile, train_dataset, val_dataset, batch_size = 1, epoch = 1000, lr = 0.0001, ft_lr = 0.000001, start=0, weight_decay = 0.005, optimizer_name="Adam", save_file=False):
@@ -22,17 +24,17 @@ def train(weights, classfile, train_dataset, val_dataset, batch_size = 1, epoch
     # apply hyper-parameters
     weight_p, bias_p, FT_weight_p, FT_bias_p, feat_weight_p, feat_bias_p = [], [], [], [], [], []
 
-    for name, p in model.named_parameters():
+    for name, mp in model.named_parameters():
         if 'FT' in name:
             if 'bias' in name:
-                FT_bias_p += [p]
+                FT_bias_p += [mp]
             else:
-                FT_weight_p += [p]
+                FT_weight_p += [mp]
         else:
             if 'bias' in name:
-                bias_p += [p]
+                bias_p += [mp]
             else:
-                weight_p += [p]
+                weight_p += [mp]
 
     # initialize optimizer and loss function
     if optimizer_name == 'SGD':
@@ -60,35 +62,39 @@ def train(weights, classfile, train_dataset, val_dataset, batch_size = 1, epoch
         dataset=train_dataset,  # torch TensorDataset format
         batch_size=batch_size,  # mini batch size
         shuffle=True,  # shuffle the dataset
-        num_workers=2,  # reading dataset by multi threads
+        num_workers=4,  # reading dataset by multi threads
         collate_fn=FolderDataSetLoader.collate_fn
     )
     # initialize maximum mAP
     max_mAP = 0
     # training
     if not os.path.exists("module"):
         os.makedirs("module")
-
-    for c_epoch in range(start, epoch):
-        # release a mini-batch data 
-        for step, (images, tags) in enumerate(loader):
-            # read images and labels
-            device_image = images.to(train_device)
-            # obtain feature maps output by the model
-            confs, locs, centers = model(device_image)  # .to(train_device)
-            # training
-            loss = loss_func(confs, locs, centers, tags, train_device)
-            optimizer.zero_grad()
-            loss.backward()
-            optimizer.step()
-            print('Epoch', c_epoch, 'Step:', step, '|train loss:%.4f' % loss)
-
-        # evaluate the performance of current model
-        mAP = return_mAP(model, val_dataset, classes)
-        print('Epoch:', c_epoch, '|mAP: %.4f' % mAP)
-        # save if better
-        if mAP >= max_mAP:
-            if save_file:
-                torch.save(model.state_dict(), save_file)
-            max_mAP = mAP
-
+    with tqdm(total=100, position=2, desc="Accuracy") as tqdm_accuracy, tqdm(total=100, position=3, desc="Recall") as tqdm_recall, tqdm(total=100, position=4, desc="Mean Average Precision") as tqdm_mAP:
+        for c_epoch in tqdm(range(start, epoch), position=0, desc="Epoch", leave=True):
+            # release a mini-batch data
+            with tqdm(enumerate(loader), total=ceil(len(train_dataset) / batch_size), unit_scale=batch_size, position=1, desc="Step", leave=True) as tdqm_enumerated_loader:
+                for step, (images, tags) in tdqm_enumerated_loader:
+                    # read images and labels
+                    device_image = images.to(train_device)
+                    # obtain feature maps output by the model
+                    confs, locs, centers = model(device_image)  # .to(train_device)
+                    # training
+                    loss = loss_func(confs, locs, centers, tags, train_device)
+                    optimizer.zero_grad()
+                    loss.backward()
+                    optimizer.step()
+                    tqdm.write('Step: %d | train loss: %.4f' % (step, loss))
+                    tdqm_enumerated_loader.set_postfix(loss = '%.4f' % loss)
+
+            # evaluate the performance of current model
+            mAP, mp, mr = return_mAP(model, val_dataset, classes)
+            tqdm_accuracy.update(mp * 100)
+            tqdm_recall.update(mr * 100)
+            tqdm_mAP.update(mAP * 100)
+            tqdm.write('Epoch: %d |mAP: %.4f' % (c_epoch, mAP))
+            # save if better
+            if mAP >= max_mAP:
+                if save_file:
+                    torch.save(model.state_dict(), save_file)
+                max_mAP = mAP

train/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "fcos-train"
-version = "0.1.6"
+version = "0.1.7"
 description = ""
 authors = ["DTheLegend <[email protected]>"]
 readme = "README.md"
@@ -9,6 +9,7 @@ packages = [{include = "fcos"}]
 [tool.poetry.dependencies]
 python = "3.10.6"
 fcos-core= {path = "../core"}
+tqdm = "^4.64.1"
 
 
 [build-system]