import datetime import os from functools import partial import numpy as np import torch import torch.backends.cudnn as cudnn import torch.distributed as dist import torch.optim as optim from torch.utils.data import DataLoader from nets.unet import Unet from nets.unet_training import get_lr_scheduler, set_optimizer_lr, weights_init from utils.callbacks import EvalCallback, LossHistory from utils.dataloader import UnetDataset, unet_dataset_collate from utils.utils import (download_weights, seed_everything, show_config, worker_init_fn) from utils.utils_fit import fit_one_epoch def train(): 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(backbone) dist.barrier() else: download_weights(backbone) model = Unet(num_classes=num_classes, pretrained=pretrained, backbone=backbone).train() if not pretrained: weights_init(model) if model_path != '': if local_rank == 0: print('Load weights {}.'.format(model_path)) 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) 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)) if local_rank == 0: time_str = datetime.datetime.strftime(datetime.datetime.now(), '%Y_%m_%d_%H_%M_%S') log_dir = os.path.join(save_dir, "loss_" + str(time_str)) loss_history = LossHistory(log_dir, model, input_shape=input_shape) else: loss_history = None if fp16: from torch.cuda.amp import GradScaler as GradScaler scaler = GradScaler() else: scaler = None model_train = model.train() 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() with open(os.path.join(VOCdevkit_path, "VOC2007/ImageSets/Segmentation/train.txt"), "r") as f: train_lines = f.readlines() with open(os.path.join(VOCdevkit_path, "VOC2007/ImageSets/Segmentation/val.txt"), "r") as f: val_lines = f.readlines() num_train = len(train_lines) num_val = len(val_lines) if True: UnFreeze_flag = False if Freeze_Train: model.freeze_backbone() batch_size = Freeze_batch_size if Freeze_Train else Unfreeze_batch_size nbs = 16 lr_limit_max = 1e-4 if optimizer_type == 'adam' else 1e-1 lr_limit_min = 1e-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 = { '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, UnFreeze_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 = UnetDataset(train_lines, input_shape, num_classes, True, VOCdevkit_path) val_dataset = UnetDataset(val_lines, input_shape, num_classes, False, VOCdevkit_path) 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=unet_dataset_collate, sampler=train_sampler, worker_init_fn=partial(worker_init_fn, rank=rank, seed=seed)) gen_val = DataLoader(val_dataset, shuffle=shuffle, batch_size=batch_size, num_workers=num_workers, pin_memory=True, drop_last=True, collate_fn=unet_dataset_collate, sampler=val_sampler, worker_init_fn=partial(worker_init_fn, rank=rank, seed=seed)) if local_rank == 0: eval_callback = EvalCallback(model, input_shape, num_classes, val_lines, VOCdevkit_path, log_dir, Cuda, \ eval_flag=eval_flag, period=eval_period) else: eval_callback = None for epoch in range(Init_Epoch, UnFreeze_Epoch): if epoch >= Freeze_Epoch and not UnFreeze_flag and Freeze_Train: batch_size = Unfreeze_batch_size nbs = 16 lr_limit_max = 1e-4 if optimizer_type == 'adam' else 1e-1 lr_limit_min = 1e-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) lr_scheduler_func = get_lr_scheduler(lr_decay_type, Init_lr_fit, Min_lr_fit, UnFreeze_Epoch) model.unfreeze_backbone() epoch_step = num_train // batch_size epoch_step_val = num_val // batch_size if epoch_step == 0 or epoch_step_val == 0: raise ValueError("数据集过小，无法继续进行训练，请扩充数据集。") if distributed: batch_size = batch_size // ngpus_per_node gen = DataLoader(train_dataset, shuffle=shuffle, batch_size=batch_size, num_workers=num_workers, pin_memory=True, drop_last=True, collate_fn=unet_dataset_collate, sampler=train_sampler, worker_init_fn=partial(worker_init_fn, rank=rank, seed=seed)) gen_val = DataLoader(val_dataset, shuffle=shuffle, batch_size=batch_size, num_workers=num_workers, pin_memory=True, drop_last=True, collate_fn=unet_dataset_collate, sampler=val_sampler, worker_init_fn=partial(worker_init_fn, rank=rank, seed=seed)) UnFreeze_flag = True if distributed: train_sampler.set_epoch(epoch) set_optimizer_lr(optimizer, lr_scheduler_func, epoch) fit_one_epoch(model_train, model, loss_history, eval_callback, optimizer, epoch, epoch_step, epoch_step_val, gen, gen_val, UnFreeze_Epoch, Cuda, dice_loss, focal_loss, cls_weights, num_classes, fp16, scaler, save_period, save_dir, local_rank) if distributed: dist.barrier() if local_rank == 0: loss_history.writer.close() if __name__ == "__main__": Cuda = True seed = 11 # 是否开启多卡 distributed = False sync_bn = False # 是否使用半精度 fp16 = True num_classes = 2 # 设置骨干网络 backbone = "vgg" pretrained = False model_path = "model_data/8414_8376.pth" input_shape = [1696, 864] # 冻结训练 Init_Epoch = 0 Freeze_Epoch = 10 Freeze_batch_size = 1 # 解冻训练 UnFreeze_Epoch = 70 Unfreeze_batch_size = 1 Freeze_Train = True # 学习率设置 Init_lr = 1e-4 Min_lr = Init_lr * 0.01 # 优化器 optimizer_type = "adam" momentum = 0.9 weight_decay = 0 lr_decay_type = 'cos' save_period = 5 save_dir = 'logs' eval_flag = True eval_period = 5 # 数据集设置 VOCdevkit_path = 'VOCdevkit' dice_loss = False focal_loss = False cls_weights = np.ones([num_classes], np.float32) num_workers = 0 seed_everything(seed) ngpus_per_node = torch.cuda.device_count() train()