添加 .gitignore 文件，更新训练和预测配置，修改模型类别，删除冗余的 readme 文件，增加测试脚本，优化学习率调度器参数

.gitignore

@@ -0,0 +1,3 @@
+*.jpg
+*.png
+__pycache__

README.md

@@ -1,5 +1,5 @@
 # 使用Pytorch简单实现SegNext，学习用
-pretrained文件夹内由官方训练的tiny版本权重文件
+pretrained文件夹内放置权重文件
 可以使用官方版本的tiny权重文件进行训练或者预测（其他版本未尝试）
 更改运行参数请使用config文件夹内的yaml文件配置
 训练移步train.py，预测移步predict.py, 有不会传的参数请移步main.py

config/lr_scheduler.yaml

@@ -0,0 +1,11 @@
+scheduler_type: LINEAR_WARMUP_THEN_POLY_SCHEDULER
+
+# total_iters=epochs * 训练的图像数量
+kwargs: |
+  {
+    "warmup_iters": 4356, 
+    "total_iters": 290400, 
+    "warmup_ratio": 0.000001, 
+    "min_lr": 0., 
+    "power": 1.
+  }

config/model.yaml

@@ -0,0 +1,45 @@
+# 网络模型的配置文件
+
+embed_dims:
+  - 3 # 输入图像的通道
+  - 32
+  - 64
+  - 160
+  - 256
+depths:
+  - 3
+  - 3
+  - 5
+  - 2
+drop_prob_of_encoder: 0.1
+expand_rations:
+  - 8
+  - 8
+  - 4
+  - 4
+
+
+
+drop_path_prob: 0.1
+drop_prob_of_decoder: 0.1
+# hidden_channels, output_channels
+channels_of_hamburger: 256
+nmf2d_config:
+  SPATIAL: 1
+  MD_S: 1
+  MD_D: 512
+  MD_R: 16
+  TRAIN_STEPS: 6
+  EVAL_STEPS: 7
+  INV_T: 1
+  ETA: 0.9
+  RAND_INIT: true
+  return_bases: false
+  device: cuda
+
+# 类别
+classes:
+  - background # 必须要
+  - 1
+  - 2
+  - 3

config/optimizer.yaml

@@ -0,0 +1,32 @@
+# 优化器的配置文件
+base_config:
+  optim_type: AdamW
+  kwargs: |
+    {
+      "lr": 0.00006,
+      "weight_decay": 0.01,
+      "betas": (0.9, 0.999)
+    }
+  
+
+
+parameters:
+    # key(block, block, norm)自定义
+    # default: decay_mult=1, lr_mult=1
+    # left必须位于第一位，表示剩余的参数，可以更改内部属性中decay_mult和lr_mult, regex_expr不起作用
+  - left:
+      regex_expr: none
+      decay_mult: 0
+      lr_mult: 1
+  - block:
+      regex_expr: ^encoder\.stages\.[0-9]\.blocks.*$
+      decay_mult: 0
+      lr_mult: 1
+  - norm:
+      regex_expr: ^encoder\.stages\.[0-9]\.norm.*$
+      decay_mult: 0
+      lr_mult: 1
+  - head:
+      regex_expr: ^decoder.*$
+      decay_mult: 1
+      lr_mult: 10

config/predict.yaml

@@ -0,0 +1,12 @@
+device: "cuda"
+
+# -1表示从不加载任何权重就进行预测
+# 0表示使用官方提供的权重进行预测
+# 1表示使用自己的权重进行预测
+mode: -1
+checkpoint:
+  - pretrained # 目录
+  - segnext_tiny_512x512_ade_160k.pth # 文件名
+# 利用正则表达式将官方model中的权重加载到模型, mode != 0时失效
+regex_expr:
+  - ^(?!decoder\.cls_seg.*)$ # ?!表示对()内的内容取反，详请见于model_utils.init_model

config/readme

@@ -0,0 +1 @@
+配置文件

config/train.yaml

@@ -0,0 +1,39 @@
+device: cuda
+batch_size: 16
+image_height: 200
+image_width: 200
+workers: 0
+epochs: 100
+
+# 每一类的占比权重，如果要让每一类的占比权重相同，为1.0即可
+weight:
+  - 1.0
+
+# 数据集存放位置
+root:
+  - dataset
+train_dir_name:
+  - train
+val_dir_name:
+  - val
+images_dir_name: images
+labels_dir_name: labels
+
+# 每训练多少个epoch就进行验证
+eval_every_n_epoch: 10
+
+# 权重保存位置
+save_path:
+  - pretrained # 目录
+  - SegNeXt # 文件名
+
+# -1表示从零开始训练网络，即不加载任何权重
+# 0表示使用官方提供的权重
+# 1表示使用自己的权重
+mode: -1
+checkpoint:
+  - pretrained # 目录
+  - segnext_tiny_512x512_ade_160k.pth # 文件名
+# 利用正则表达式将官方model中的权重加载到模型, mode != 0时失效
+regex_expr:
+  - ^(?!decoder\.cls_seg.*)$ # ?!表示对()内的内容取反, 详请见于model_utils.init_model

learning_rate_scheduler.py

@@ -187,7 +187,7 @@ class LinearWarmupThenPolyScheduler:
         total_iters: 总训练步数
         min_lr: 最低学习率
     """
-    def __init__(self, optimizer, warmup_iters=1500, total_iters=2000, warmup_ratio=1e-6, min_lr=0., power=1.):
+    def __init__(self, optimizer, warmup_iters=1500, total_iters=20000, warmup_ratio=1e-6, min_lr=0., power=1.):
         self.optimizer = optimizer
         self.current_iters = 0
         self.warmup_iters = warmup_iters

test.ipynb

@@ -0,0 +1,110 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from PIL import Image\n",
+    "import numpy as np"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "(200, 200, 3)\n",
+      "[[[226 226 226]\n",
+      "  [227 227 227]\n",
+      "  [225 225 225]\n",
+      "  ...\n",
+      "  [153 153 153]\n",
+      "  [153 153 153]\n",
+      "  [152 152 152]]\n",
+      "\n",
+      " [[227 227 227]\n",
+      "  [222 222 222]\n",
+      "  [221 221 221]\n",
+      "  ...\n",
+      "  [151 151 151]\n",
+      "  [150 150 150]\n",
+      "  [149 149 149]]\n",
+      "\n",
+      " [[227 227 227]\n",
+      "  [226 226 226]\n",
+      "  [221 221 221]\n",
+      "  ...\n",
+      "  [153 153 153]\n",
+      "  [152 152 152]\n",
+      "  [153 153 153]]\n",
+      "\n",
+      " ...\n",
+      "\n",
+      " [[234 234 234]\n",
+      "  [229 229 229]\n",
+      "  [229 229 229]\n",
+      "  ...\n",
+      "  [160 160 160]\n",
+      "  [159 159 159]\n",
+      "  [159 159 159]]\n",
+      "\n",
+      " [[231 231 231]\n",
+      "  [233 233 233]\n",
+      "  [230 230 230]\n",
+      "  ...\n",
+      "  [159 159 159]\n",
+      "  [160 160 160]\n",
+      "  [159 159 159]]\n",
+      "\n",
+      " [[230 230 230]\n",
+      "  [236 236 236]\n",
+      "  [234 234 234]\n",
+      "  ...\n",
+      "  [155 155 155]\n",
+      "  [156 156 156]\n",
+      "  [154 154 154]]]\n"
+     ]
+    }
+   ],
+   "source": [
+    "# 读取图片\n",
+    "image_path = 'dataset/train/images/000229.jpg'\n",
+    "image = Image.open(image_path)\n",
+    "\n",
+    "# 将图片转换为 NumPy 数组\n",
+    "image_array = np.array(image)\n",
+    "\n",
+    "# 显示图片的像素值\n",
+    "print(image_array.shape)\n",
+    "print(image_array)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "torch",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.9"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

train.py

@@ -260,7 +260,8 @@ def train(
             f=f"{os.path.sep.join(train_config['save_path'])}_val_Iou{100 * best_val_model['avg_iou']:.3f}_{datetime.strftime(datetime.now(), '%Y%m%d%H%M%S')}.pth"
         )
 
-
+def count_parameters(model):
+    return sum(p.numel() for p in model.parameters() if p.requires_grad)
 
 
 
@@ -273,6 +274,11 @@ if __name__ == "__main__":
         net=net,
         optimizer=optimizer
     )
+
+    # 计算并打印网络的参数量
+    # num_params = count_parameters(net)
+    # print(f"网络的参数量: {num_params}")
+
     train(
         net=net,
         optimizer=optimizer,