391 lines
16 KiB
Python
391 lines
16 KiB
Python
|
#!/usr/bin/env python
|
|||
|
|
# -*- coding: utf-8 -*-
|
|||
|
|
"""
|
|||
|
|
@project:
|
|||
|
|
@File : train_module
|
|||
|
|
@Author : root
|
|||
|
|
@create_time : 2022/9/13 22:57
|
|||
|
|
训练文件模板
|
|||
|
|
"""
|
|||
|
|
import os
|
|||
|
|
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
|
|||
|
|
import sys
|
|||
|
|
root_dir='/home/qiqq/q3dl/codeR/applicationprojectR/'
|
|||
|
|
sys.path.append(root_dir)
|
|||
|
|
from taihuyuan_roof.utils.index_manager import addLogger,addTensorboard
|
|||
|
|
from taihuyuan_roof.utils.util import getLogger
|
|||
|
|
|
|||
|
|
|
|||
|
|
from taihuyuan_roof.manet.model.resunet import resUnetpamcarb
|
|||
|
|
from taihuyuan_roof.manet.train_val_test.evaluate_train import Evalue
|
|||
|
|
|
|||
|
|
|
|||
|
|
import argparse
|
|||
|
|
import logging
|
|||
|
|
|
|||
|
|
from pathlib import Path
|
|||
|
|
from torch import nn
|
|||
|
|
import torch.nn.functional as F
|
|||
|
|
from torch import optim
|
|||
|
|
from torch.utils.data import DataLoader
|
|||
|
|
from tqdm import tqdm
|
|||
|
|
from tensorboardX import SummaryWriter
|
|||
|
|
|
|||
|
|
import time
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
"由于服务器用不了wandb所以这里删除了用wandb记录的代码"
|
|||
|
|
import os
|
|||
|
|
import numpy as np
|
|||
|
|
import random
|
|||
|
|
import torch
|
|||
|
|
seed = 123
|
|||
|
|
random.seed(seed)
|
|||
|
|
np.random.seed(seed)
|
|||
|
|
torch.manual_seed(seed) # 为CPU设置随机种子
|
|||
|
|
torch.cuda.manual_seed(seed) # 为当前GPU设置随机种子
|
|||
|
|
|
|||
|
|
|
|||
|
|
# 加了tensorboard记录
|
|||
|
|
dir_checkpoint = Path('./checkpoints/')
|
|||
|
|
tensorboarddir= Path('./TensorboardLog/')
|
|||
|
|
logsavedir = Path('./logs/')
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
#在train上的权重[0.5, 2.5, 2.5]
|
|||
|
|
def criterion(inputs, target, loss_weight=None, num_classes: int = 2, dice: bool = True, ignore_index: int = -100):
|
|||
|
|
|
|||
|
|
l_loss = F.cross_entropy(inputs, target, ignore_index=ignore_index, weight=loss_weight)
|
|||
|
|
|
|||
|
|
|
|||
|
|
return l_loss
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
##cityscpes 我直接改成20类了19类原来的+1类背景（原来255位置的）
|
|||
|
|
def train_net(net,
|
|||
|
|
device,
|
|||
|
|
resume=False,
|
|||
|
|
isPretrain=False,
|
|||
|
|
epochs: int = 5,
|
|||
|
|
train_batch_size: int = 2,
|
|||
|
|
val_batch_size: int = 2,
|
|||
|
|
learning_rate: float = 1e-3,
|
|||
|
|
save_checkpoint: bool = True,
|
|||
|
|
ignoreindex: int = 100,
|
|||
|
|
num_class = 21,
|
|||
|
|
useDice=False,
|
|||
|
|
valIndex=[],
|
|||
|
|
|
|||
|
|
|
|||
|
|
):
|
|||
|
|
logger = getLogger(logsavedir)
|
|||
|
|
# 1. Create dataset
|
|||
|
|
# myself数据加载
|
|||
|
|
|
|||
|
|
from taihuyuan_roof.dataloaders.datasets import def_taihuyuan
|
|||
|
|
parsertrain = argparse.ArgumentParser()
|
|||
|
|
argstrain = parsertrain.parse_args()
|
|||
|
|
argstrain.resize = (512, 512)
|
|||
|
|
argstrain.crop_size = 480 # 如果不开始数据增强，这个两个也是没用的
|
|||
|
|
argstrain.flip_prob = 0.5
|
|||
|
|
|
|||
|
|
parserval = argparse.ArgumentParser()
|
|||
|
|
argsval = parserval.parse_args()
|
|||
|
|
argsval.resize = (512, 512)
|
|||
|
|
argsval.crop_size = 480 # 在验证的时候没有用 crop和filp 只是为了不报错
|
|||
|
|
argsval.flip_prob = 0.5
|
|||
|
|
imgshape_base = None
|
|||
|
|
|
|||
|
|
train_d = def_taihuyuan.datasets_rooftaihuyuan(argstrain, split='trainbat', isAug=True)
|
|||
|
|
val_d = def_taihuyuan.datasets_rooftaihuyuan(argsval, split='valbat', )
|
|||
|
|
|
|||
|
|
n_val = val_d.__len__()
|
|||
|
|
n_train = train_d.__len__()
|
|||
|
|
|
|||
|
|
train_loader = DataLoader(train_d, batch_size=train_batch_size, shuffle=True, num_workers=4, pin_memory=False,
|
|||
|
|
drop_last=False)
|
|||
|
|
val_loader = DataLoader(val_d, shuffle=False, num_workers=4, pin_memory=False, batch_size=val_batch_size,
|
|||
|
|
drop_last=False)
|
|||
|
|
|
|||
|
|
|
|||
|
|
workTensorboard_dir = os.path.join(tensorboarddir,
|
|||
|
|
time.strftime("%Y-%m-%d-%H.%M", time.localtime())) # 日志文件写入目录
|
|||
|
|
if not os.path.exists(workTensorboard_dir):
|
|||
|
|
os.makedirs(workTensorboard_dir)
|
|||
|
|
|
|||
|
|
workcheckpoint_dir = os.path.join(dir_checkpoint,
|
|||
|
|
time.strftime("%Y-%m-%d-%H.%M", time.localtime())) # 日志文件写入目录
|
|||
|
|
if not os.path.exists(workcheckpoint_dir):
|
|||
|
|
os.makedirs(workcheckpoint_dir)
|
|||
|
|
|
|||
|
|
writer = SummaryWriter(workTensorboard_dir)
|
|||
|
|
description="优化器:Adam,学习率:LinoPolyScheduler,\n" \
|
|||
|
|
"其他策略:无数据增强，无类加权，归一化方式/255\n" \
|
|||
|
|
"res50的预训练\n" \
|
|||
|
|
"网络：res+Unetpam+carb(MANET) \n" \
|
|||
|
|
"32倍数下次采样。 \n" \
|
|||
|
|
"数据集：太湖源roof\n " \
|
|||
|
|
"保存最好的权重的指标是roof的iou,\n"
|
|||
|
|
|
|||
|
|
|
|||
|
|
logger.info(f'''Starting training:
|
|||
|
|
"Description:" {description}
|
|||
|
|
Epochs: {epochs}
|
|||
|
|
train_Batch size: {train_batch_size}
|
|||
|
|
val_Batch size: {val_batch_size}
|
|||
|
|
Learning rate: {learning_rate}
|
|||
|
|
Training size: {n_train}
|
|||
|
|
Validation size: {n_val}
|
|||
|
|
Checkpoints: {save_checkpoint}
|
|||
|
|
Device: {device.type}
|
|||
|
|
|
|||
|
|
|
|||
|
|
''')
|
|||
|
|
net = net.to(device)
|
|||
|
|
from taihuyuan_roof.schedulers.polyscheduler import LinoPolyScheduler
|
|||
|
|
|
|||
|
|
# 4. Set up the optimizer, the loss, the learning rate scheduler and the loss scaling for AMP
|
|||
|
|
optimizer = optim.Adam(net.parameters(), lr=learning_rate)
|
|||
|
|
# optimizer = optim.SGD(net.parameters(),lr=learning_rate,momentum=0.9,weight_decay=1e-4)
|
|||
|
|
scheduler = LinoPolyScheduler(optimizer, epochs=epochs ,steps_per_epoch=len(train_loader),min_lr=0, )
|
|||
|
|
# scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.98, )
|
|||
|
|
# scheduler = optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer,T_0=5,T_mult=2,eta_min=1e-6)
|
|||
|
|
|
|||
|
|
# we = np.array([1,3,16], np.float32)
|
|||
|
|
# we = torch.from_numpy(we).to(device)
|
|||
|
|
|
|||
|
|
|
|||
|
|
start_epoch = 1
|
|||
|
|
# #是否有预训练
|
|||
|
|
# if isPretrain:
|
|||
|
|
# model_dict = net.state_dict()
|
|||
|
|
# pretrained_dict = torch.load(isPretrain, 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)
|
|||
|
|
# net.load_state_dict(model_dict)
|
|||
|
|
# print("train中的预训练加载成功")
|
|||
|
|
# logger.info(f'Model loaded from {isPretrain}')
|
|||
|
|
# else:#没有预训练使用pyroch的一些权重初始化方法
|
|||
|
|
# strweight='normal'
|
|||
|
|
# weights_init(net,init_type=strweight)
|
|||
|
|
# logging.info(f'没有预训练权重，{strweight}权重初始化完成')
|
|||
|
|
# # #是否使用断点训练
|
|||
|
|
if resume:
|
|||
|
|
path_checkpoint = resume # 断点路径
|
|||
|
|
checkpoint = torch.load(path_checkpoint) # 加载断点
|
|||
|
|
net.load_state_dict(checkpoint['net']) # 加载模型可学习参数
|
|||
|
|
# 加载优化器参数
|
|||
|
|
start_epoch = checkpoint['epoch'] +1 # 设置开始的epoch
|
|||
|
|
scheduler.load_state_dict(checkpoint['lr_schedule']) # 加载lr_scheduler
|
|||
|
|
print("重启成功")
|
|||
|
|
else:
|
|||
|
|
start_epoch= 1
|
|||
|
|
|
|||
|
|
|
|||
|
|
# 5. Begin training
|
|||
|
|
epochs_score = [] # 记录每个epoh的miou
|
|||
|
|
best_miou = 0.0 # 记录最好的那个
|
|||
|
|
earlyStop_count=0
|
|||
|
|
|
|||
|
|
time1 = time.time()
|
|||
|
|
start = time.strftime('%Y-%m-%d %H:%M:%S',time.localtime(time1))
|
|||
|
|
print("训练{}个epoch的开始时间为:{}".format(epochs,start))
|
|||
|
|
|
|||
|
|
logger.info("训练{}个epoch的开始时间为:{}".format(epochs,start))
|
|||
|
|
|
|||
|
|
log = addLogger(logger)
|
|||
|
|
tenb = addTensorboard(writer)
|
|||
|
|
|
|||
|
|
# early_stopping = EarlyStopping(patience=100)
|
|||
|
|
|
|||
|
|
for epoch in range(start_epoch, epochs + 1):
|
|||
|
|
logger.info("train|epoch:{epoch}\t".format(epoch=epoch))
|
|||
|
|
current_miou = 0.0
|
|||
|
|
total_train_loss = 0
|
|||
|
|
net.train()
|
|||
|
|
print('Start Train')
|
|||
|
|
time1_1 = time.time()
|
|||
|
|
|
|||
|
|
with tqdm(total=len(train_loader), desc=f'Epoch {epoch}/{epochs}', unit='batch') as pbar:
|
|||
|
|
for iteration, batch in enumerate(train_loader): # batchsize=2 一共1487.5个batch
|
|||
|
|
images = batch['image']
|
|||
|
|
true_masks = batch['label']
|
|||
|
|
# true_masks =true_masks-1#这一条是针对有的8位彩色图索引从1开始
|
|||
|
|
images = images.to(device=device, dtype=torch.float32)
|
|||
|
|
true_masks = true_masks.to(device=device, dtype=torch.long)
|
|||
|
|
|
|||
|
|
# out0,out1 = net(images, use128aux=True,state="train")
|
|||
|
|
out = net(images)
|
|||
|
|
loss = 0.0
|
|||
|
|
if useDice:
|
|||
|
|
loss = criterion(out, true_masks, loss_weight=None, num_classes=num_class, dice= True, ignore_index = ignoreindex)
|
|||
|
|
|
|||
|
|
else:
|
|||
|
|
|
|||
|
|
if isinstance(out, list) or isinstance(out, tuple): # 可能有多个输出（这里把辅助解码头的也输出的所以是多个）
|
|||
|
|
|
|||
|
|
loss0 = criterion(out[0], true_masks, loss_weight=None, num_classes=num_class, dice=False,
|
|||
|
|
ignore_index=ignoreindex)
|
|||
|
|
|
|||
|
|
loss1 = criterion(out[1], true_masks, loss_weight=None, num_classes=num_class, dice=False,
|
|||
|
|
ignore_index=ignoreindex)
|
|||
|
|
|
|||
|
|
loss=loss0+0.4*loss1
|
|||
|
|
else:
|
|||
|
|
loss0 = criterion(out, true_masks, loss_weight=None, num_classes=num_class, dice=False,
|
|||
|
|
ignore_index=ignoreindex)
|
|||
|
|
loss = loss0
|
|||
|
|
|
|||
|
|
|
|||
|
|
'''1.loss 2.梯度清零，3.反向传播。backward 4optomizer更新.'''
|
|||
|
|
|
|||
|
|
optimizer.zero_grad()
|
|||
|
|
loss.backward()
|
|||
|
|
optimizer.step()
|
|||
|
|
pbar.update() # 用来记录一次更新多少的
|
|||
|
|
total_train_loss += loss.item()
|
|||
|
|
|
|||
|
|
pbar.set_postfix(**{'loss (batch)': loss.item()})
|
|||
|
|
|
|||
|
|
|
|||
|
|
writer.add_scalar("train_total_loss", total_train_loss / (iteration + 1), epoch)
|
|||
|
|
time1_2 = time.time()
|
|||
|
|
one_epoch_time = round((time1_2-time1_1)/60,3) #按照分钟来算
|
|||
|
|
logger.info("训练一个epoch大概:{} min".format(one_epoch_time))
|
|||
|
|
print("Finish Train")
|
|||
|
|
print('start Validation')
|
|||
|
|
isevalue = True
|
|||
|
|
|
|||
|
|
if isevalue == True:
|
|||
|
|
val_score =0
|
|||
|
|
current_lr= optimizer.param_groups[0]['lr'] #当前的学习率
|
|||
|
|
evalueModel=Evalue(net,val_loader,device,num_class,isResize=imgshape_base,ignore_index=ignoreindex,weight=None)
|
|||
|
|
for i in valIndex:
|
|||
|
|
#这个地方其实是默认了只有IouMiouP可能会出现数组
|
|||
|
|
if i =="IouMiouP":
|
|||
|
|
indextring = [ "acc_global", "acc", "iu","precion","recall","f1","miou","meanf1","roof_iou","roof_f1"]
|
|||
|
|
result=eval('evalueModel.'+'evalue_'+i)()
|
|||
|
|
for j in range(len(indextring)):
|
|||
|
|
log.logger_index(result[j],indextring[j])
|
|||
|
|
if indextring[j] == "acc_global" or indextring[j] == "miou" or indextring[j] == "meanf1" or indextring[j] =="roof_iou"or indextring[j] =="roof_f1":
|
|||
|
|
if indextring[j] == "roof_iou":
|
|||
|
|
val_score = result[j]
|
|||
|
|
current_miou=val_score
|
|||
|
|
epochs_score.append(val_score)
|
|||
|
|
tenb.writer_singleindex(indextring[j],index=result[j],epoch=epoch)
|
|||
|
|
|
|||
|
|
else:
|
|||
|
|
tenb.writer_classindex(indextring[j],classindexs=result[j],epoch=epoch)
|
|||
|
|
else:
|
|||
|
|
result = eval('evalueModel.' + 'evalue_' + i)()
|
|||
|
|
|
|||
|
|
log.logger_index(result, i)
|
|||
|
|
tenb.writer_singleindex(i, index=result, epoch=epoch)
|
|||
|
|
|
|||
|
|
writer.add_scalar("lr", current_lr, epoch)
|
|||
|
|
logger.info("当前epoch{}的lr:{}".format(epoch, current_lr))
|
|||
|
|
scheduler.step()
|
|||
|
|
# 保存最好的miou和最新的
|
|||
|
|
if save_checkpoint:
|
|||
|
|
checkpoint = {
|
|||
|
|
"net": net.state_dict(),
|
|||
|
|
'optimizer': optimizer.state_dict(),
|
|||
|
|
"epoch": epoch,
|
|||
|
|
'lr_schedule': scheduler.state_dict()
|
|||
|
|
}
|
|||
|
|
torch.save(checkpoint, os.path.join(workcheckpoint_dir, 'last.pth')) # 保存最新的
|
|||
|
|
# 保存最好的
|
|||
|
|
if current_miou >= best_miou:
|
|||
|
|
best_miou = current_miou
|
|||
|
|
torch.save(checkpoint, os.path.join(workcheckpoint_dir, 'best.pth'))
|
|||
|
|
writer.add_scalar('best_epoch_index', epochs_score.index(max(epochs_score)) + 1, epoch) #记录bestepoch的变化情况
|
|||
|
|
# 这个地方是按照迭代来调整学习率的
|
|||
|
|
print('Finish Validation')
|
|||
|
|
# 达到早停止条件时，early_stop会被置为True
|
|||
|
|
# early_stopping(val_score)
|
|||
|
|
# if early_stopping.early_stop:
|
|||
|
|
# print("Early stopping")
|
|||
|
|
# time2 = time.time()
|
|||
|
|
# end = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(time2))
|
|||
|
|
# logger.info("训练提前结束:训练{}个epoch的结束时间为:{}".format(epoch, end))
|
|||
|
|
# break
|
|||
|
|
|
|||
|
|
time2 = time.time()
|
|||
|
|
end = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(time2))
|
|||
|
|
print("训练结束:训练{}个epoch的结束时间为:{}".format(epochs, end))
|
|||
|
|
logger.info("训练结束:训练{}个epoch的结束时间为:{}".format(epochs, end))
|
|||
|
|
writer.close()
|
|||
|
|
logger.info("训练完成")
|
|||
|
|
|
|||
|
|
|
|||
|
|
def get_args():
|
|||
|
|
parser = argparse.ArgumentParser(description='Train the UNet on images and target masks')
|
|||
|
|
parser.add_argument('--epochs', '-e', metavar='E', type=int, default=200, help='Number of epochs')
|
|||
|
|
parser.add_argument('--train_batch-size', '-tb', dest='train_batch_size', metavar='TB', type=int, default=8,
|
|||
|
|
help='Train_Batch size')
|
|||
|
|
parser.add_argument('--val_batch-size', '-vb', dest='val_batch_size', metavar='VB', type=int, default=1,
|
|||
|
|
help='Val_Batch size')
|
|||
|
|
parser.add_argument('--learning-rate', '-l', metavar='LR', type=float, default=1e-3,
|
|||
|
|
help='Learning rate', dest='lr')
|
|||
|
|
parser.add_argument('--load', '-f', type=str,
|
|||
|
|
default="/home/qiqq/q3dl/code/pretrain_weight/pretrained/resnet/resnet50-0676ba61.pth",
|
|||
|
|
help='Load model from a .pth file') # 有没有预训练。。
|
|||
|
|
# parser.add_argument('--load', '-f', type=str,
|
|||
|
|
# default="/home/qiqq/q3dl/code/applicationproject/taihuyuan_roof/manet/train_val_test/checkpoints/2023-09-20-14.07/best.pth",
|
|||
|
|
# help='Load model from a .pth file') # 有没有预训练。。
|
|||
|
|
parser.add_argument('--ignore_index', '-i', type=int, dest='ignore_index', default=255,
|
|||
|
|
help='ignore index defult 100')
|
|||
|
|
parser.add_argument('--origin_shape', action='store_true', default=(512,512), help='原始输入尺寸')
|
|||
|
|
parser.add_argument('--resume', '-r', type=str, default="", help='is use Resume')
|
|||
|
|
parser.add_argument('--useDice', '-ud', type=str, default=False, help='训练的时候是否使用dice')
|
|||
|
|
parser.add_argument('--valIndex', '-vI', type=str, default=["Valloss","IouMiouP"], help='评价指标要使用哪些,注意IouMiouP= acc_global, acc, iu,precion,recall,f1,miou,并且建议旨在2分类的时候用dice否则会出错')
|
|||
|
|
parser.add_argument('--classes', '-c', type=int, default=2, help='Number of classes')
|
|||
|
|
|
|||
|
|
return parser.parse_args()
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == '__main__':
|
|||
|
|
|
|||
|
|
# 有点乱啊，
|
|||
|
|
args = get_args() # 管着非网络的参数的
|
|||
|
|
|
|||
|
|
logging.basicConfig(level=logging.INFO, format='%(levelname)s: %(message)s')
|
|||
|
|
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
|||
|
|
logging.info(f'Using device {device}')
|
|||
|
|
|
|||
|
|
# net igin_shape,pretrain_backbone=args.load)
|
|||
|
|
net= resUnetpamcarb()
|
|||
|
|
try:
|
|||
|
|
train_net(net=net,
|
|||
|
|
device = device,
|
|||
|
|
resume=args.resume,
|
|||
|
|
epochs=args.epochs,
|
|||
|
|
isPretrain=args.load,
|
|||
|
|
train_batch_size=args.train_batch_size,
|
|||
|
|
val_batch_size=args.val_batch_size,
|
|||
|
|
learning_rate=args.lr,
|
|||
|
|
ignoreindex=args.ignore_index,
|
|||
|
|
num_class=args.classes,
|
|||
|
|
useDice=args.useDice,
|
|||
|
|
valIndex=args.valIndex
|
|||
|
|
|
|||
|
|
)
|
|||
|
|
except KeyboardInterrupt:
|
|||
|
|
torch.save(net.state_dict(), 'INTERRUPTED.pth')
|
|||
|
|
logging.info('Saved interrupt')
|
|||
|
|
raise
|