lichen/ieemoo-ai-contrast
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

多机并行计算

Browse Source
This commit is contained in:
lee
2025-08-14 10:09:54 +08:00
parent bc896fc688
commit 99a204ee22
18 changed files with 105 additions and 55 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
.gitignore vendored
View File

@ -8,4 +8,5 @@ loss/
checkpoints/
search_library/
quant_imgs/
electronic_imgs/
README.md

10
configs/compare.yml
View File

@ -15,8 +15,8 @@ base:
# 模型配置
models:
backbone: 'resnet50'
channel_ratio: 1.0
backbone: 'resnet18'
channel_ratio: 0.75
# 训练参数
training:
@ -31,9 +31,9 @@ training:
weight_decay: 0.0005 # 权重衰减
scheduler: "step" # 学习率调度器（可选：cosine/cosine_warm/step/None）
num_workers: 32 # 数据加载线程数
checkpoints: "./checkpoints/resnet50_electornic_20250807/" # 模型保存目录
checkpoints: "./checkpoints/resnet18_pdd_test/" # 模型保存目录
restore: false
restore_model: "./checkpoints/resnet18_20250717_scale=0.75_nosub/best.pth" # 模型恢复路径
restore_model: "./checkpoints/resnet50_electornic_20250807/best.pth" # 模型恢复路径
cosine_t_0: 10 # 初始周期长度
cosine_t_mult: 1 # 周期长度倍率
cosine_eta_min: 0.00001 # 最小学习率
@ -70,3 +70,5 @@ logging:
distributed:
enabled: false # 是否启用分布式训练
backend: "nccl" # 分布式后端（nccl/gloo）
node_rank: 0 # 节点编号
node_num: 1 # 共计几个节点 一般几台机器就有几个节点

3
configs/pic_pic_similar.yml
View File

@ -15,7 +15,8 @@ base:
models:
backbone: 'resnet18'
channel_ratio: 0.75
model_path: "../checkpoints/resnet18_1009/best.pth"
model_path: "../checkpoints/resnet18_20250715_scale=0.75_sub/best.pth"
# model_path: "../checkpoints/resnet18_1009/best.pth"
heatmap:
feature_layer: "layer4"

3
configs/similar_analysis.yml
View File

@ -15,7 +15,8 @@ base:
models:
backbone: 'resnet18'
channel_ratio: 0.75
model_path: "../checkpoints/resnet18_1009/best.pth"
# model_path: "../checkpoints/resnet18_1009/best.pth"
model_path: "../checkpoints/resnet18_20250715_scale=0.75_sub/best.pth"
heatmap:
feature_layer: "layer4"

10
configs/test.yml
View File

@ -13,8 +13,8 @@ base:
# 模型配置
models:
backbone: 'resnet18'
channel_ratio: 0.75
model_path: "checkpoints/resnet18_1009/best.pth"
channel_ratio: 1.0
model_path: "checkpoints/resnet18_electornic_20250806/best.pth"
#resnet18_20250715_scale=0.75_sub
#resnet18_20250718_scale=0.75_nosub
half: false # 是否启用半精度测试（fp16）
@ -24,11 +24,11 @@ models:
data:
test_batch_size: 128 # 训练批次大小
num_workers: 32 # 数据加载线程数
test_dir: "../data_center/contrast_data/v1/extra" # 验证数据集根目录
test_dir: "../data_center/electornic/v1/val" # 验证数据集根目录
test_group_json: "../data_center/contrast_learning/model_test_data/test/inner_group_pairs.json"
test_list: "../data_center/contrast_data/v1/extra_cross_same.txt"
test_list: "../data_center/electornic/v1/cross_same.txt"
group_test: false
save_image_joint: true
save_image_joint: false
image_joint_pth: "./joint_images"
transform:

9
configs/transform.yml
View File

@ -10,15 +10,16 @@ base:
embedding_size: 256 # 特征维度
pin_memory: true # 是否启用pin_memory
distributed: true # 是否启用分布式训练
dataset: "./dataset_electornic.txt" # 数据集名称
# 模型配置
models:
backbone: 'resnet18'
backbone: 'resnet101'
channel_ratio: 1.0
model_path: "../checkpoints/resnet18_1009/best.pth"
onnx_model: "../checkpoints/resnet18_3399_sancheng/best.onnx"
rknn_model: "../checkpoints/resnet18_3399_sancheng/best_rknn2.3.2_RK3566.rknn"
model_path: "../checkpoints/resnet101_electornic_20250807/best.pth"
onnx_model: "../checkpoints/resnet101_electornic_20250807/best.onnx"
rknn_model: "../checkpoints/resnet101_electornic_20250807/resnet101_electornic.rknn"
rknn_batch_size: 1
# 日志与监控

9
configs/utils.py
View File

@ -1,4 +1,4 @@
from model import (resnet18, resnet34, resnet50, mobilevit_s, MobileNetV3_Small, MobileNetV3_Large, mobilenet_v1,
from model import (resnet18, resnet34, resnet50, resnet101, mobilevit_s, MobileNetV3_Small, MobileNetV3_Large, mobilenet_v1,
PPLCNET_x1_0, PPLCNET_x0_5, PPLCNET_x2_5)
from timm.models import vit_base_patch16_224 as vit_base_16
from model.metric import ArcFace, CosFace
@ -13,9 +13,10 @@ class trainer_tools:
def get_backbone(self):
backbone_mapping = {
'resnet18': lambda: resnet18(scale=self.conf['models']['channel_ratio']),
'resnet34': lambda: resnet34(scale=self.conf['models']['channel_ratio']),
'resnet50': lambda: resnet50(scale=self.conf['models']['channel_ratio']),
'resnet18': lambda: resnet18(scale=self.conf['models']['channel_ratio'], pretrained=True),
'resnet34': lambda: resnet34(scale=self.conf['models']['channel_ratio'], pretrained=True),
'resnet50': lambda: resnet50(scale=self.conf['models']['channel_ratio'], pretrained=True),
'resnet101': lambda: resnet101(scale=self.conf['models']['channel_ratio'], pretrained=True),
'mobilevit_s': lambda: mobilevit_s(),
'mobilenetv3_small': lambda: MobileNetV3_Small(),
'PPLCNET_x1_0': lambda: PPLCNET_x1_0(),

12
configs/write_feature.yml
View File

@ -12,9 +12,9 @@ base:
# 模型配置
models:
backbone: 'resnet18'
channel_ratio: 0.75
checkpoints: "../checkpoints/resnet18_20250715_scale=0.75_sub/best.pth"
backbone: 'resnet101'
channel_ratio: 1.0
checkpoints: "../checkpoints/resnet101_electornic_20250807/best.pth"
# 数据配置
data:
@ -22,7 +22,7 @@ data:
test_batch_size: 128 # 验证批次大小
num_workers: 32 # 数据加载线程数
half: true # 是否启用半精度数据
img_dirs_path: "/home/lc/data_center/baseStlib/pic/stlib_base" # base标准库图片存储路径
img_dirs_path: "/shareData/completed_data/scatter_data/electronic_scale/base/total" # base标准库图片存储路径
# img_dirs_path: "/home/lc/contrast_nettest/data/feature_json"
xlsx_pth: false # 过滤商品， 默认None不进行过滤
@ -41,8 +41,8 @@ logging:
checkpoint_interval: 30 # 检查点保存间隔（epoch）
save:
json_bin: "../search_library/yunhedian_05-09.json" # 保存整个json文件
json_bin: "../search_library/resnet101_electronic.json" # 保存整个json文件
json_path: "/home/lc/data_center/baseStlib/feature_json/stlib_base_resnet18_sub" # 保存单个json文件路径
error_barcodes: "error_barcodes.txt"
barcodes_statistics: "../search_library/barcodes_statistics.txt"
create_single_json: true # 是否保存单个json文件
create_single_json: false # 是否保存单个json文件

2
model/__init__.py
View File

@ -4,7 +4,7 @@ from .mobilevit import mobilevit_s
from .metric import ArcFace, CosFace
from .loss import FocalLoss
from .resbam import resnet
from .resnet_pre import resnet18, resnet34, resnet50, resnet14, CustomResNet18
from .resnet_pre import resnet18, resnet34, resnet50, resnet101, resnet152,resnet14, CustomResNet18
from .mobilenet_v2 import mobilenet_v2
from .mobilenet_v3 import MobileNetV3_Small, MobileNetV3_Large
# from .mobilenet_v1 import mobilenet_v1

6
model/resnet_pre.py
View File

@ -368,7 +368,7 @@ def resnet18(pretrained=True, progress=True, **kwargs):
**kwargs)
def resnet34(pretrained=False, progress=True, **kwargs):
def resnet34(pretrained=True, progress=True, **kwargs):
r"""ResNet-34 model from
`"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_
@ -380,7 +380,7 @@ def resnet34(pretrained=False, progress=True, **kwargs):
**kwargs)
def resnet50(pretrained=False, progress=True, **kwargs):
def resnet50(pretrained=True, progress=True, **kwargs):
r"""ResNet-50 model from
`"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_
@ -392,7 +392,7 @@ def resnet50(pretrained=False, progress=True, **kwargs):
**kwargs)
def resnet101(pretrained=False, progress=True, **kwargs):
def resnet101(pretrained=True, progress=True, **kwargs):
r"""ResNet-101 model from
`"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_

2
test_ori.py
View File

@ -17,7 +17,7 @@ from configs import trainer_tools
import yaml
from datetime import datetime
with open('./configs/test.yml', 'r') as f:
with open('../configs/test.yml', 'r') as f:
conf = yaml.load(f, Loader=yaml.FullLoader)
# Constants from config

23
tools/dataset_electornic.txt Normal file
View File

@ -0,0 +1,23 @@
../electronic_imgs/0.jpg
../electronic_imgs/1.jpg
../electronic_imgs/2.jpg
../electronic_imgs/3.jpg
../electronic_imgs/4.jpg
../electronic_imgs/5.jpg
../electronic_imgs/6.jpg
../electronic_imgs/7.jpg
../electronic_imgs/8.jpg
../electronic_imgs/9.jpg
../electronic_imgs/10.jpg
../electronic_imgs/11.jpg
../electronic_imgs/12.jpg
../electronic_imgs/13.jpg
../electronic_imgs/14.jpg
../electronic_imgs/15.jpg
../electronic_imgs/16.jpg
../electronic_imgs/17.jpg
../electronic_imgs/18.jpg
../electronic_imgs/19.jpg
../electronic_imgs/20.jpg
../electronic_imgs/21.jpg
../electronic_imgs/22.jpg

4
tools/getpairs.py
View File

@ -203,11 +203,11 @@ class PairGenerator:
if __name__ == "__main__":
original_path = '/home/lc/data_center/contrast_data/v1/extra'
original_path = '/home/lc/data_center/electornic/v1/val'
parent_dir = str(Path(original_path).parent)
generator = PairGenerator(original_path)
# Example usage:
pairs = generator.get_pairs(original_path,
output_txt=os.sep.join([parent_dir, 'extra_cross_same.txt'])) # Individual pairs
output_txt=os.sep.join([parent_dir, 'cross_same.txt'])) # Individual pairs
# groups = generator.get_group_pairs('val') # Group pairs

5
tools/model_rknn_transform.py
View File

@ -122,8 +122,9 @@ if __name__ == '__main__':
# Build model
print('--> Building model')
ret = rknn.build(do_quantization=False, # True
dataset='./dataset.txt',
ret = rknn.build(do_quantization=True, # True
# dataset='./dataset.txt',
dataset=conf['base']['dataset'],
rknn_batch_size=conf['models']['rknn_batch_size'])
# ret = rknn.build(do_quantization=False, dataset='./dataset.txt')
if ret != 0:

6
tools/picdir_to_picdir_similar.py
View File

@ -237,6 +237,6 @@ def get_histogram(data, label=None):
if __name__ == '__main__':
# picTopic_matrix = picDirSimilarAnalysis()
# picTopic_matrix.get_group_similarity_matrix('/home/lc/data_center/image_analysis/pic_pic_similar_maxtrix')
read_result_txt()
picTopic_matrix = picDirSimilarAnalysis()
picTopic_matrix.get_group_similarity_matrix('/home/lc/data_center/image_analysis/pic_pic_similar_maxtrix')
# read_result_txt()

2
tools/similar_analysis.py
View File

@ -22,7 +22,7 @@ class SimilarAnalysis:
"""初始化模型和度量方法"""
tr_tools = trainer_tools(conf)
backbone_mapping = tr_tools.get_backbone()
print('model_path {}'.format(conf['models']['model_path']))
if conf['models']['backbone'] in backbone_mapping:
model = backbone_mapping[conf['models']['backbone']]()
else:

38
tools/write_feature_json.py
View File

@ -4,7 +4,7 @@ import logging
import numpy as np
from typing import Dict, List, Optional, Tuple
from tools.dataset import get_transform
from model import resnet18
from model import resnet18, resnet34, resnet50, resnet101
import torch
from PIL import Image
import pandas as pd
@ -50,7 +50,16 @@ class FeatureExtractor:
raise FileNotFoundError(f"Model weights file not found: {model_path}")
# Initialize model
model = resnet18(scale=self.conf['models']['channel_ratio']).to(self.conf['base']['device'])
if conf['models']['backbone'] == 'resnet18':
model = resnet18(scale=self.conf['models']['channel_ratio']).to(self.conf['base']['device'])
elif conf['models']['backbone'] == 'resnet34':
model = resnet34(scale=self.conf['models']['channel_ratio']).to(self.conf['base']['device'])
elif conf['models']['backbone'] == 'resnet50':
model = resnet50(scale=self.conf['models']['channel_ratio']).to(self.conf['base']['device'])
elif conf['models']['backbone'] == 'resnet101':
model = resnet101(scale=self.conf['models']['channel_ratio']).to(self.conf['base']['device'])
else:
print("不支持的模型: {}".format(conf['models']['backbone']))
# Handle multi-GPU case
if conf['base']['distributed']:
@ -168,7 +177,7 @@ class FeatureExtractor:
# Validate input directory
if not os.path.isdir(folder):
raise ValueError(f"Invalid directory: {folder}")
i = 0
# Process each barcode directory
for root, dirs, files in tqdm(os.walk(folder), desc="Scanning directories"):
if not dirs: # Leaf directory (contains images)
@ -180,14 +189,16 @@ class FeatureExtractor:
ori_barcode = basename
barcode = basename
# Apply filter if provided
i += 1
print(ori_barcode, i)
if filter and ori_barcode not in filter:
continue
elif len(ori_barcode) > 13 or len(ori_barcode) < 8:
logger.warning(f"Skipping invalid barcode {ori_barcode}")
with open(conf['save']['error_barcodes'], 'a') as f:
f.write(ori_barcode + '\n')
f.close()
continue
# elif len(ori_barcode) > 13 or len(ori_barcode) < 8: # barcode筛选长度
# logger.warning(f"Skipping invalid barcode {ori_barcode}")
# with open(conf['save']['error_barcodes'], 'a') as f:
# f.write(ori_barcode + '\n')
# f.close()
# continue
# Process image files
if files:
@ -299,7 +310,8 @@ class FeatureExtractor:
dicts['value'] = truncated_imgs_list
if create_single_json:
# json_path = os.path.join("./search_library/v8021_overseas/", str(barcode_list[i]) + '.json')
json_path = os.path.join(self.conf['save']['json_path'], str(barcode_list[i]) + '.json')
json_path = os.path.join(self.conf['save']['json_path'],
str(barcode_list[i]) + '.json')
with open(json_path, 'w') as json_file:
json.dump(dicts, json_file)
else:
@ -317,8 +329,10 @@ class FeatureExtractor:
with open(conf['save']['barcodes_statistics'], 'w', encoding='utf-8') as f:
for barcode in os.listdir(pth):
print("barcode length >> {}".format(len(barcode)))
if len(barcode) > 13 or len(barcode) < 8:
continue
# if len(barcode) > 13 or len(barcode) < 8: # barcode筛选长度
# continue
if filter is not None:
f.writelines(barcode + '\n')
if barcode in filter:

15
train_compare.py
View File

@ -266,11 +266,15 @@ def main():
if distributed:
# 分布式训练：使用mp.spawn启动多个进程
world_size = torch.cuda.device_count()
local_size = torch.cuda.device_count()
world_size = int(conf['distributed']['node_num'])*local_size
mp.spawn(
run_training,
args=(world_size, conf),
nprocs=world_size,
args=(conf['distributed']['node_rank'],
local_size,
world_size,
conf),
nprocs=local_size,
join=True
)
else:
@ -279,11 +283,12 @@ def main():
run_training_loop(components)
def run_training(rank, world_size, conf):
def run_training(local_rank, node_rank, local_size, world_size, conf):
"""实际执行训练的函数，供mp.spawn调用"""
# 初始化分布式环境
rank = local_rank + node_rank * local_size
os.environ['RANK'] = str(rank)
os.environ['WORLD_SIZE'] = str(world_size)
os.environ['WORLD_SIZE'] = str(local_size)
os.environ['MASTER_ADDR'] = 'localhost'
os.environ['MASTER_PORT'] = '12355'