310 lines
10 KiB
Python
310 lines
10 KiB
Python
import os
|
||
import os.path as osp
|
||
|
||
import torch
|
||
import torch.nn as nn
|
||
import torch.optim as optim
|
||
from tqdm import tqdm
|
||
from torch.nn.parallel import DistributedDataParallel as DDP
|
||
import torch.distributed as dist
|
||
import torch.multiprocessing as mp
|
||
import torch.utils.data.distributed
|
||
|
||
from model.loss import FocalLoss
|
||
from tools.dataset import load_data
|
||
import matplotlib.pyplot as plt
|
||
from configs import trainer_tools
|
||
import yaml
|
||
from datetime import datetime
|
||
|
||
|
||
def load_config(config_path='configs/scatter.yml'):
|
||
"""加载配置文件."""
|
||
with open(config_path, 'r') as f:
|
||
return yaml.load(f, Loader=yaml.FullLoader)
|
||
|
||
|
||
# 加载配置
|
||
conf = load_config()
|
||
|
||
|
||
# 数据加载封装
|
||
def load_datasets():
|
||
"""加载训练和验证数据集，并为分布式训练创建 DistributedSampler."""
|
||
train_dataset, class_num = load_data(training=True, cfg=conf)
|
||
val_dataset, _ = load_data(training=False, cfg=conf)
|
||
|
||
if conf['base']['distributed']:
|
||
train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
|
||
val_sampler = torch.utils.data.distributed.DistributedSampler(val_dataset)
|
||
else:
|
||
train_sampler = None
|
||
val_sampler = None
|
||
|
||
train_dataloader = torch.utils.data.DataLoader(
|
||
train_dataset,
|
||
batch_size=conf['data']['batch_size'],
|
||
shuffle=(train_sampler is None),
|
||
num_workers=conf['data']['num_workers'],
|
||
pin_memory=conf['data']['pin_memory'],
|
||
sampler=train_sampler
|
||
)
|
||
|
||
val_dataloader = torch.utils.data.DataLoader(
|
||
val_dataset,
|
||
batch_size=conf['data']['batch_size'],
|
||
shuffle=False,
|
||
num_workers=conf['data']['num_workers'],
|
||
pin_memory=conf['data']['pin_memory'],
|
||
sampler=val_sampler
|
||
)
|
||
|
||
return train_dataloader, val_dataloader, class_num
|
||
|
||
|
||
# 加载数据集
|
||
train_dataloader, val_dataloader, class_num = load_datasets()
|
||
|
||
tr_tools = trainer_tools(conf)
|
||
backbone_mapping = tr_tools.get_backbone()
|
||
metric_mapping = tr_tools.get_metric(class_num)
|
||
|
||
|
||
# 设备管理封装
|
||
def get_device(device_config=None):
|
||
"""根据配置返回设备（CPU/GPU）。在分布式环境下初始化进程组."""
|
||
if device_config is None:
|
||
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
||
else:
|
||
device = torch.device(device_config)
|
||
|
||
if conf['base']['distributed']:
|
||
dist.init_process_group(backend='nccl')
|
||
|
||
return device
|
||
|
||
|
||
# 获取设备
|
||
device = get_device(conf['base']['device'])
|
||
|
||
# 模型初始化
|
||
if conf['models']['backbone'] in backbone_mapping:
|
||
model = backbone_mapping[conf['models']['backbone']]().to(device)
|
||
else:
|
||
raise ValueError('不支持该模型: {}'.format({conf['models']['backbone']}))
|
||
|
||
if conf['training']['metric'] in metric_mapping:
|
||
metric = metric_mapping[conf['training']['metric']]().to(device)
|
||
else:
|
||
raise ValueError('不支持的metric类型: {}'.format(conf['training']['metric']))
|
||
|
||
rank = 0
|
||
if torch.cuda.device_count() > 1 and conf['base']['distributed']:
|
||
print("Let's use", torch.cuda.device_count(), "GPUs!")
|
||
dist.barrier()
|
||
model = DDP(model, device_ids=[rank])
|
||
metric = DDP(metric, device_ids=[rank])
|
||
|
||
# Training Setup
|
||
def initialize_components():
|
||
# 封装模型、损失函数和优化器的初始化
|
||
if conf['training']['loss'] == 'focal_loss':
|
||
criterion = FocalLoss(gamma=2)
|
||
else:
|
||
criterion = nn.CrossEntropyLoss()
|
||
|
||
optimizer_mapping = tr_tools.get_optimizer(model, metric)
|
||
if conf['training']['optimizer'] in optimizer_mapping:
|
||
optimizer = optimizer_mapping[conf['training']['optimizer']]()
|
||
scheduler_mapping = tr_tools.get_scheduler(optimizer)
|
||
scheduler = scheduler_mapping[conf['training']['scheduler']]()
|
||
print('使用{}优化器 使用{}调度器'.format(conf['training']['optimizer'],
|
||
conf['training']['scheduler']))
|
||
return criterion, optimizer, scheduler
|
||
else:
|
||
raise ValueError('不支持的优化器类型: {}'.format(conf['training']['optimizer']))
|
||
|
||
|
||
# 初始化组件
|
||
criterion, optimizer, scheduler = initialize_components()
|
||
|
||
# Checkpoints Setup
|
||
checkpoints = conf['training']['checkpoints']
|
||
os.makedirs(checkpoints, exist_ok=True)
|
||
|
||
|
||
def train_epoch(model, dataloader, optimizer, criterion, device):
|
||
model.train()
|
||
train_loss = 0
|
||
for data, labels in tqdm(dataloader, desc="Training", ascii=True, total=len(dataloader)):
|
||
data, labels = data.to(device), labels.to(device)
|
||
embeddings = model(data).to(device)
|
||
if not conf['training']['metric'] == 'softmax':
|
||
thetas = metric(embeddings, labels)
|
||
else:
|
||
thetas = metric(embeddings)
|
||
loss = criterion(thetas, labels)
|
||
optimizer.zero_grad()
|
||
loss.backward()
|
||
optimizer.step()
|
||
train_loss += loss.item()
|
||
return train_loss / len(dataloader)
|
||
|
||
|
||
def validate_epoch(model, dataloader, criterion, device):
|
||
model.eval()
|
||
val_loss = 0
|
||
with torch.no_grad():
|
||
for data, labels in tqdm(dataloader, desc="Validation", ascii=True, total=len(dataloader)):
|
||
data, labels = data.to(device), labels.to(device)
|
||
embeddings = model(data).to(device)
|
||
if not conf['training']['metric'] == 'softmax':
|
||
thetas = metric(embeddings, labels)
|
||
else:
|
||
thetas = metric(embeddings)
|
||
loss = criterion(thetas, labels)
|
||
val_loss += loss.item()
|
||
return val_loss / len(dataloader)
|
||
|
||
|
||
def save_model(model, path, distributed):
|
||
if distributed and torch.cuda.device_count() > 1:
|
||
if dist.get_rank() == 0:
|
||
torch.save(model.module.state_dict(), path)
|
||
else:
|
||
torch.save(model.state_dict(), path)
|
||
|
||
|
||
def write_log(log_info, log_dir):
|
||
with open(log_dir, 'a') as f:
|
||
f.write(log_info + '\n')
|
||
|
||
|
||
def plot_losses(epochs, train_losses, val_losses, save_path):
|
||
plt.plot(epochs, train_losses, color='blue', label='Train Loss')
|
||
plt.plot(epochs, val_losses, color='red', label='Validation Loss')
|
||
plt.xlabel('Epochs')
|
||
plt.ylabel('Loss')
|
||
plt.legend()
|
||
plt.savefig(save_path)
|
||
plt.close()
|
||
|
||
|
||
# 模型恢复封装
|
||
def restore_model(model, device):
|
||
if conf['training']['restore']:
|
||
print('load pretrain model: {}'.format(conf['training']['restore_model']))
|
||
model.load_state_dict(torch.load(conf['training']['restore_model'], map_location=device))
|
||
return model
|
||
|
||
|
||
# 日志和学习率记录封装
|
||
def log_and_print(e, train_loss_avg, val_loss_avg, current_lr, log_dir):
|
||
if conf['base']['distributed'] and dist.get_rank() != 0:
|
||
return
|
||
|
||
log_info = ("[{:%Y-%m-%d %H:%M:%S}] Epoch {}/{}, train_loss: {}, val_loss: {} lr:{}"
|
||
.format(datetime.now(),
|
||
e,
|
||
conf['training']['epochs'],
|
||
train_loss_avg,
|
||
val_loss_avg,
|
||
current_lr))
|
||
print(log_info)
|
||
write_log(log_info, log_dir)
|
||
print("第%d个epoch的学习率：%f" % (e, current_lr))
|
||
|
||
|
||
# 模型评估与保存封装
|
||
def evaluate_and_save(val_loss_avg, best_loss, model, checkpoints, distributed):
|
||
if val_loss_avg < best_loss:
|
||
best_path = osp.join(checkpoints, 'best.pth')
|
||
save_model(model, best_path, distributed)
|
||
best_loss = val_loss_avg
|
||
return best_loss
|
||
|
||
|
||
def run_training(rank, world_size, conf):
|
||
"""在指定 rank 上运行训练 loop。
|
||
|
||
Args:
|
||
rank: 当前进程的索引。
|
||
world_size: 进程总数。
|
||
conf: 配置字典。
|
||
"""
|
||
dist.init_process_group(backend='nccl', rank=rank, world_size=world_size)
|
||
device = torch.device('cuda', rank)
|
||
|
||
# 数据加载器和模型等需要重新初始化以确保每个进程独立工作
|
||
train_dataloader, val_dataloader, class_num = load_datasets()
|
||
tr_tools = trainer_tools(conf)
|
||
backbone_mapping = tr_tools.get_backbone()
|
||
metric_mapping = tr_tools.get_metric(class_num)
|
||
|
||
# 模型初始化
|
||
if conf['models']['backbone'] in backbone_mapping:
|
||
model = backbone_mapping[conf['models']['backbone']]().to(device)
|
||
else:
|
||
raise ValueError('不支持该模型: {}'.format({conf['models']['backbone']}))
|
||
|
||
if conf['training']['metric'] in metric_mapping:
|
||
metric = metric_mapping[conf['training']['metric']]().to(device)
|
||
else:
|
||
raise ValueError('不支持的metric类型: {}'.format(conf['training']['metric']))
|
||
|
||
model = DDP(model, device_ids=[rank])
|
||
metric = DDP(metric, device_ids=[rank])
|
||
|
||
# 初始化组件
|
||
criterion, optimizer, scheduler = initialize_components()
|
||
|
||
# 恢复模型（如果需要）
|
||
model = restore_model(model, device)
|
||
|
||
train_losses = []
|
||
val_losses = []
|
||
epochs = []
|
||
temp_loss = 1000
|
||
log_dir = osp.join(conf['logging']['logging_dir'])
|
||
|
||
for e in range(conf['training']['epochs']):
|
||
train_loss_avg = train_epoch(model, train_dataloader, optimizer, criterion, device)
|
||
train_losses.append(train_loss_avg)
|
||
|
||
val_loss_avg = validate_epoch(model, val_dataloader, criterion, device)
|
||
val_losses.append(val_loss_avg)
|
||
|
||
temp_loss = evaluate_and_save(val_loss_avg, temp_loss, model, conf['training']['checkpoints'], True)
|
||
|
||
scheduler.step()
|
||
current_lr = optimizer.param_groups[0]['lr']
|
||
log_and_print(e, train_loss_avg, val_loss_avg, current_lr, log_dir)
|
||
epochs.append(e)
|
||
|
||
last_path = osp.join(conf['training']['checkpoints'], 'last.pth')
|
||
save_model(model, last_path, True)
|
||
plot_losses(epochs, train_losses, val_losses, 'loss/mobilenetv3Large_2250_0316.png')
|
||
|
||
dist.destroy_process_group()
|
||
|
||
|
||
def main():
|
||
world_size = torch.cuda.device_count()
|
||
mp.spawn(
|
||
run_training,
|
||
args=(world_size, conf),
|
||
nprocs=world_size,
|
||
join=True,
|
||
)
|
||
|
||
|
||
if __name__ == '__main__':
|
||
print('backbone>{} '.format(conf['models']['backbone']),
|
||
'metric>{} '.format(conf['training']['metric']),
|
||
'checkpoints>{} '.format(conf['training']['checkpoints']),
|
||
)
|
||
if conf['base']['distributed']:
|
||
main()
|
||
else:
|
||
run_training(rank=0, world_size=1, conf=conf)
|