# ----------------------------------------------------#
# 是否使用Cuda
# 没有GPU可以设置成False
# ----------------------------------------------------#
Cuda = True
# ---------------------------------------------------------------------#
# distributed 用于指定是否使用单机多卡分布式运行
# 终端指令仅支持Ubuntu。CUDA_VISIBLE_DEVICES用于在Ubuntu下指定显卡。
# Windows系统下默认使用DP模式调用所有显卡，不支持DDP。
# DP模式：
# 设置 distributed = False
# 在终端中输入 CUDA_VISIBLE_DEVICES=0,1 python train.py
# DDP模式：
# 设置 distributed = True
# 在终端中输入 CUDA_VISIBLE_DEVICES=0,1 python -m torch.distributed.launch --nproc_per_node=2 train.py
# ---------------------------------------------------------------------#
distributed = False
# ---------------------------------------------------------------------#
# sync_bn 是否使用sync_bn，DDP模式多卡可用
# ---------------------------------------------------------------------#
sync_bn = False
# ---------------------------------------------------------------------#
# fp16 是否使用混合精度训练
# 可减少约一半的显存、需要pytorch1.7.1以上
# ---------------------------------------------------------------------#
fp16 = False
# ----------------------------------------------------#
# 数据集存放的路径
# ----------------------------------------------------#
dataset_path = "datasets"
# ----------------------------------------------------#
# 输入图像的大小，默认为105,105,3
# ----------------------------------------------------#
input_shape = [105, 105]
# ----------------------------------------------------#
# 当训练Omniglot数据集时设置为False
# 当训练自己的数据集时设置为True
#
# 训练自己的数据和Omniglot数据格式不一样。
# 详情可看README.md
# ----------------------------------------------------#
train_own_data = True
# -------------------------------#
# 用于指定是否使用VGG预训练权重
# 有两种获取方式
# 1、利用百度网盘下载后放入
# ./model_data/
# 2、直接运行自动下载
# -------------------------------#
pretrained = True
# ----------------------------------------------------------------------------------------------------------------------------#
# 权值文件的下载请看README，可以通过网盘下载。模型的 预训练权重 对不同数据集是通用的，因为特征是通用的。
# 模型的 预训练权重 比较重要的部分是 主干特征提取网络的权值部分，用于进行特征提取。
# 预训练权重对于99%的情况都必须要用，不用的话主干部分的权值太过随机，特征提取效果不明显，网络训练的结果也不会好
#
# 如果训练过程中存在中断训练的操作，可以将model_path设置成logs文件夹下的权值文件，将已经训练了一部分的权值再次载入。
# 同时修改下方的 冻结阶段 或者 解冻阶段 的参数，来保证模型epoch的连续性。
#
# 当model_path = ''的时候不加载整个模型的权值。
#
# 此处使用的是整个模型的权重，因此是在train.py进行加载的，pretrain不影响此处的权值加载。
# 如果想要让模型从主干的预训练权值开始训练，则设置model_path = ''，pretrain = True，此时仅加载主干。
# 如果想要让模型从0开始训练，则设置model_path = ''，pretrain = Fasle，此时从0开始训练。
# 一般来讲，从0开始训练效果会很差，因为权值太过随机，特征提取效果不明显。
#
# 网络一般不从0开始训练，至少会使用主干部分的权值，有些论文提到可以不用预训练，主要原因是他们 数据集较大 且 调参能力优秀。
# 如果一定要训练网络的主干部分，可以了解imagenet数据集，首先训练分类模型，分类模型的 主干部分 和该模型通用，基于此进行训练。
# ----------------------------------------------------------------------------------------------------------------------------#
model_path = "./logs/best_epoch_weights.pth"
# ----------------------------------------------------------------------------------------------------------------------------#
# 显存不足与数据集大小无关，提示显存不足请调小batch_size。
# 受到BatchNorm层影响，不能为1。
#
# 在此提供若干参数设置建议，各位训练者根据自己的需求进行灵活调整：
# （一）从预训练权重开始训练：
# Adam：
# Init_Epoch = 0，Epoch = 100，optimizer_type = 'adam'，Init_lr = 1e-3，weight_decay = 0。
# SGD：
# Init_Epoch = 0，Epoch = 100，optimizer_type = 'sgd'，Init_lr = 1e-2，weight_decay = 5e-4。
# 其中：UnFreeze_Epoch可以在100-300之间调整。
# （二）batch_size的设置：
# 在显卡能够接受的范围内，以大为好。显存不足与数据集大小无关，提示显存不足（OOM或者CUDA out of memory）请调小batch_size。
# 受到BatchNorm层影响，batch_size最小为2，不能为1。
# 正常情况下Freeze_batch_size建议为Unfreeze_batch_size的1-2倍。不建议设置的差距过大，因为关系到学习率的自动调整。
# ----------------------------------------------------------------------------------------------------------------------------#
# ------------------------------------------------------#
# 训练参数
# Init_Epoch 模型当前开始的训练世代
# Epoch 模型总共训练的epoch
# batch_size 每次输入的图片数量
# ------------------------------------------------------#
Init_Epoch = 3
Epoch = 100
batch_size = 64
# ------------------------------------------------------------------#
# 其它训练参数：学习率、优化器、学习率下降有关
# ------------------------------------------------------------------#
# ------------------------------------------------------------------#
# Init_lr 模型的最大学习率
# 当使用Adam优化器时建议设置 Init_lr=1e-3
# 当使用SGD优化器时建议设置 Init_lr=1e-2
# Min_lr 模型的最小学习率，默认为最大学习率的0.01
# ------------------------------------------------------------------#
Init_lr = 1e-2
Min_lr = Init_lr * 0.01
# ------------------------------------------------------------------#
# optimizer_type 使用到的优化器种类，可选的有adam、sgd
# 当使用Adam优化器时建议设置 Init_lr=1e-3
# 当使用SGD优化器时建议设置 Init_lr=1e-2
# momentum 优化器内部使用到的momentum参数
# weight_decay 权值衰减，可防止过拟合
# adam会导致weight_decay错误，使用adam时建议设置为0。
# ------------------------------------------------------------------#
optimizer_type = "sgd"
momentum = 0.9
weight_decay = 5e-4
# ------------------------------------------------------------------#
# lr_decay_type 使用到的学习率下降方式，可选的有'step'、'cos'
# ------------------------------------------------------------------#
lr_decay_type = 'cos'
# ------------------------------------------------------------------#
# save_period 多少个epoch保存一次权值
# ------------------------------------------------------------------#
save_period = 10
# ------------------------------------------------------------------#
# save_dir 权值与日志文件保存的文件夹
# ------------------------------------------------------------------#
save_dir = 'logs'
# ------------------------------------------------------------------#
# num_workers 用于设置是否使用多线程读取数据，1代表关闭多线程
# 开启后会加快数据读取速度，但是会占用更多内存
# 在IO为瓶颈的时候再开启多线程，即GPU运算速度远大于读取图片的速度。
# ------------------------------------------------------------------#
num_workers = 8
# ------------------------------------------------------#
# 设置用到的显卡
# ------------------------------------------------------#
ngpus_per_node = torch.cuda.device_count()
if distributed:
dist.init_process_group(backend="nccl")
local_rank = int(os.environ["LOCAL_RANK"])
rank = int(os.environ["RANK"])
device = torch.device("cuda", local_rank)
if local_rank == 0:
print(f"[{os.getpid()}] (rank = {rank}, local_rank = {local_rank}) training...")
print("Gpu Device Count : ", ngpus_per_node)
else:
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
local_rank = 0
rank = 0
if pretrained:
if distributed:
if local_rank == 0:
download_weights("vgg16")
dist.barrier()
else:
download_weights("vgg16")
model = Siamese(input_shape, pretrained)
if model_path != '':
# ------------------------------------------------------#
# 权值文件请看README，百度网盘下载
# ------------------------------------------------------#
if local_rank == 0:
print('Load weights {}.'.format(model_path))
# ------------------------------------------------------#
# 根据预训练权重的Key和模型的Key进行加载
# ------------------------------------------------------#
model_dict = model.state_dict()
pretrained_dict = torch.load(model_path, map_location=device)
load_key, no_load_key, temp_dict = [], [], {}
for k, v in pretrained_dict.items():
if k in model_dict.keys() and np.shape(model_dict[k]) == np.shape(v):
temp_dict[k] = v
load_key.append(k)
else:
no_load_key.append(k)
model_dict.update(temp_dict)
model.load_state_dict(model_dict)
# ------------------------------------------------------#
# 显示没有匹配上的Key
# ------------------------------------------------------#
if local_rank == 0:
print("\nSuccessful Load Key:", str(load_key)[:500], "……\nSuccessful Load Key Num:", len(load_key))
print("\nFail To Load Key:", str(no_load_key)[:500], "……\nFail To Load Key num:", len(no_load_key))
print("\n\033[1;33;44m温馨提示，head部分没有载入是正常现象，Backbone部分没有载入是错误的。\033[0m")
# ----------------------#
# 获得损失函数
# ----------------------#
loss = nn.BCEWithLogitsLoss()
# ----------------------#
# 记录Loss
# ----------------------#
if local_rank == 0:
loss_history = LossHistory(save_dir, model, input_shape=input_shape)
else:
loss_history = None
# ------------------------------------------------------------------#
# torch 1.2不支持amp，建议使用torch 1.7.1及以上正确使用fp16
# 因此torch1.2这里显示"could not be resolve"
# ------------------------------------------------------------------#
if fp16:
from torch.cuda.amp import GradScaler as GradScaler
scaler = GradScaler()
else:
scaler = None
model_train = model.train()
# ----------------------------#
# 多卡同步Bn
# ----------------------------#
if sync_bn and ngpus_per_node > 1 and distributed:
model_train = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model_train)
elif sync_bn:
print("Sync_bn is not support in one gpu or not distributed.")
if Cuda:
if distributed:
# ----------------------------#
# 多卡平行运行
# ----------------------------#
model_train = model_train.cuda(local_rank)
model_train = torch.nn.parallel.DistributedDataParallel(model_train, device_ids=[local_rank],
find_unused_parameters=True)
else:
model_train = torch.nn.DataParallel(model)
cudnn.benchmark = True
model_train = model_train.cuda()
# ----------------------------------------------------#
# 训练集和验证集的比例。
# ----------------------------------------------------#
train_ratio = 0.9
train_lines, train_labels, val_lines, val_labels = load_dataset(dataset_path, train_own_data, train_ratio)
num_train = len(train_lines)
num_val = len(val_lines)
if local_rank == 0:
show_config(
model_path=model_path, input_shape=input_shape,
Init_Epoch=Init_Epoch, Epoch=Epoch, batch_size=batch_size,
Init_lr=Init_lr, Min_lr=Min_lr, optimizer_type=optimizer_type, momentum=momentum,
lr_decay_type=lr_decay_type,
save_period=save_period, save_dir=save_dir, num_workers=num_workers, num_train=num_train, num_val=num_val
)
# ---------------------------------------------------------#
# 总训练世代指的是遍历全部数据的总次数
# 总训练步长指的是梯度下降的总次数
# 每个训练世代包含若干训练步长，每个训练步长进行一次梯度下降。
# 此处仅建议最低训练世代，上不封顶，计算时只考虑了解冻部分
# ----------------------------------------------------------#
wanted_step = 3e4 if optimizer_type == "sgd" else 1e4
total_step = num_train // batch_size * Epoch
if total_step <= wanted_step:
wanted_epoch = wanted_step // (num_train // batch_size) + 1
print("\n\033[1;33;44m[Warning] 使用%s优化器时，建议将训练总步长设置到%d以上。\033[0m" % (
optimizer_type, wanted_step))
print(
"\033[1;33;44m[Warning] 本次运行的总训练数据量为%d，batch_size为%d，共训练%d个Epoch，计算出总训练步长为%d。\033[0m" % (
num_train, batch_size, Epoch, total_step))
print("\033[1;33;44m[Warning] 由于总训练步长为%d，小于建议总步长%d，建议设置总世代为%d。\033[0m" % (
total_step, wanted_step, wanted_epoch))
# -------------------------------------------------------------#
# 训练分为两个阶段，两阶段初始的学习率不同，手动调节了学习率
# 显存不足与数据集大小无关，提示显存不足请调小batch_size。
# -------------------------------------------------------------#
if True:
# -------------------------------------------------------------------#
# 判断当前batch_size，自适应调整学习率
# -------------------------------------------------------------------#
nbs = 64
lr_limit_max = 1e-3 if optimizer_type == 'adam' else 1e-1
lr_limit_min = 3e-4 if optimizer_type == 'adam' else 5e-4
Init_lr_fit = min(max(batch_size / nbs * Init_lr, lr_limit_min), lr_limit_max)
Min_lr_fit = min(max(batch_size / nbs * Min_lr, lr_limit_min * 1e-2), lr_limit_max * 1e-2)
# ---------------------------------------#
# 根据optimizer_type选择优化器
# ---------------------------------------#
optimizer = {
'adam': optim.Adam(model.parameters(), Init_lr_fit, betas=(momentum, 0.999), weight_decay=weight_decay),
'sgd': optim.SGD(model.parameters(), Init_lr_fit, momentum=momentum, nesterov=True,
weight_decay=weight_decay)
}[optimizer_type]
# ---------------------------------------#
# 获得学习率下降的公式
# ---------------------------------------#
lr_scheduler_func = get_lr_scheduler(lr_decay_type, Init_lr_fit, Min_lr_fit, Epoch)
# ---------------------------------------#
# 判断每一个世代的长度
# ---------------------------------------#
epoch_step = num_train // batch_size
epoch_step_val = num_val // batch_size
if epoch_step == 0 or epoch_step_val == 0:
raise ValueError("数据集过小，无法继续进行训练，请扩充数据集。")
train_dataset = SiameseDataset(input_shape, train_lines, train_labels, True)
val_dataset = SiameseDataset(input_shape, val_lines, val_labels, False)
if distributed:
train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset, shuffle=True, )
val_sampler = torch.utils.data.distributed.DistributedSampler(val_dataset, shuffle=False, )
batch_size = batch_size // ngpus_per_node
shuffle = False
else:
train_sampler = None
val_sampler = None
shuffle = True
gen = DataLoader(train_dataset, shuffle=shuffle, batch_size=batch_size, num_workers=num_workers,
pin_memory=True,
drop_last=True, collate_fn=dataset_collate, sampler=train_sampler)
gen_val = DataLoader(val_dataset, shuffle=shuffle, batch_size=batch_size, num_workers=num_workers,
pin_memory=True,
drop_last=True, collate_fn=dataset_collate, sampler=val_sampler)
for epoch in range(Init_Epoch, Epoch):
if distributed:
train_sampler.set_epoch(epoch)
set_optimizer_lr(optimizer, lr_scheduler_func, epoch)
fit_one_epoch(model_train, model, loss, loss_history, optimizer, epoch, epoch_step, epoch_step_val, gen,
gen_val, Epoch, Cuda, fp16, scaler, save_period, save_dir, local_rank)
if local_rank == 0:
loss_history.writer.close()
React
Typescript
Django
Laravel
Facebook
Microsoft
Google
Alibaba
D3
Tencent