[hrnet]最终运行程序

[hrnet]最终运行程序

仅作为记录，大佬请跳过

项目：传送门

HRNet-Image-Classification

文章目录

• /tools/train.py
• /tools/valid.py
• /lib/config/default.py
• /experiments/cls_hrnet_w18_sgd_lr5e-2_wd1e-4_bs32_x100.yaml

/tools/train.py

# ------------------------------------------------------------------------------
# Written by Bin Xiao (Bin.Xiao@microsoft)
# Modified by Ke Sun (sunk@mail.ustc.edu)
# ------------------------------------------------------------------------------from __future__ import absolute_import
from __future__ import division
from __future__ import print_functionimport argparse
import os
os.environ['CUDA_VISIBLE_DEVICES']='1/'
import pprint
import shutil
import sysimport torch
parallel
import torch.backends.cudnn as cudnn
import torch.optim
import torch.utils.data
import torch.utils.data.distributed
import torchvision.datasets as datasets
ansforms as transforms
from tensorboardX import SummaryWriterimport _init_paths
import models
from config import config
from config import update_config
from core.function import train
from core.function import validate
delsummary import get_model_summary
from utils.utils import get_optimizer
from utils.utils import save_checkpoint
from utils.utils import create_loggerdef parse_args():parser = argparse.ArgumentParser(description='Train classification network')parser.add_argument('--cfg',help='experiment configure file name',required=True,type=str)parser.add_argument('--modelDir',help='model directory',type=str,default='')parser.add_argument('--logDir',help='log directory',type=str,default='')parser.add_argument('--dataDir',help='data directory',type=str,default='')parser.add_argument('--testModel',help='testModel',type=str,default='')args = parser.parse_args()update_config(config, args)return argsdef main():args = parse_args()logger, final_output_dir, tb_log_dir = create_logger(config, args.cfg, 'train')logger.info(pprint.pformat(args))logger.info(pprint.pformat(config))# cudnn related settingcudnn.benchmark = config.CUDNN.BENCHMARKtorch.backends.cudnn.deterministic = config.CUDNN.DETERMINISTICtorch.abled = config.CUDNN.ENABLEDmodel = eval('models.'+config.MODEL.NAME+'.get_cls_net')(config)dump_input = torch.rand((1, 3, config.MODEL.IMAGE_SIZE[1], config.MODEL.IMAGE_SIZE[0]))logger.info(get_model_summary(model, dump_input))# copy model filethis_dir = os.path.dirname(__file__)models_dst_dir = os.path.join(final_output_dir, 'models')if ists(models_dst_dir):(models_dst_pytree(os.path.join(this_dir, '../lib/models'), models_dst_dir)writer_dict = {'writer': SummaryWriter(log_dir=tb_log_dir),'train_global_steps': 0,'valid_global_steps': 0,}gpus = list(config.GPUS)
#     model = DataParallel(model, device_ids=gpus).cuda()device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")model&#(device)# define loss function (criterion) and optimizercriterion = CrossEntropyLoss().cuda()optimizer = get_optimizer(config, model)best_perf = 0.0best_model = Falselast_epoch = config.TRAIN.BEGIN_EPOCHif config.TRAIN.RESUME:model_state_file = os.path.join(final_output_dir,'checkpoint.pth.tar')if os.path.isfile(model_state_file):checkpoint = torch.load(model_state_file)last_epoch = checkpoint['epoch']best_perf = checkpoint['perf']model.load_state_dict(checkpoint['state_dict'])# dule.load_state_dict(checkpoint['state_dict'])optimizer.load_state_dict(checkpoint['optimizer'])logger.info("=> loaded checkpoint (epoch {})".format(checkpoint['epoch']))best_model = Trueif isinstance(config.TRAIN.LR_STEP, list):lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, config.TRAIN.LR_STEP, config.TRAIN.LR_FACTOR,last_epoch-1)else:lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, config.TRAIN.LR_STEP, config.TRAIN.LR_FACTOR,last_epoch-1)# Data loading codetraindir = os.path.join(config.DATASET.ROOT, config.DATASET.TRAIN_SET)valdir = os.path.join(config.DATASET.ROOT, config.DATASET.TEST_SET)normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],std=[0.229, 0.224, 0.225])train_dataset = datasets.ImageFolder(traindir,transforms.Compose([transforms.RandomResizedCrop(config.MODEL.IMAGE_SIZE[0]),transforms.RandomHorizontalFlip(),transforms.ToTensor(),normalize,]))#以下为增加的代码，上面几行是原有的代码#print(train_dataset.classes)  #根据分的文件夹的名字来确定的类别with open(&#","w") as f1:for classname in train_dataset.classes:f1.write(classname + "n")#print(train_dataset.class_to_idx) #按顺序为这些类别定义索引为with open(&#", "w") as f2:for key, value in train_dataset.class_to_idx.items():f2.write(str(key) + " " + str(value) + 'n')#print(train_dataset.imgs) #返回从所有文件夹中得到的图片的路径以及其类别train_loader = torch.utils.data.DataLoader(train_dataset,batch_size=100,#config.TRAIN.BATCH_SIZE_PER_GPU*len(gpus),shuffle=True,num_workers=config.WORKERS,pin_memory=True)valid_loader = torch.utils.data.DataLoader(datasets.ImageFolder(valdir, transforms.Compose([transforms.Resize(int(config.MODEL.IMAGE_SIZE[0] / 0.875)),transforms.CenterCrop(config.MODEL.IMAGE_SIZE[0]),transforms.ToTensor(),normalize,])),batch_size=100,#config.TEST.BATCH_SIZE_PER_GPU*len(gpus),shuffle=False,num_workers=config.WORKERS,pin_memory=True)for epoch in range(last_epoch, config.TRAIN.END_EPOCH):lr_scheduler.step()# train for one epochtrain(config, train_loader, model, criterion, optimizer, epoch,final_output_dir, tb_log_dir, writer_dict)# evaluate on validation setperf_indicator = validate(config, valid_loader, model, criterion,final_output_dir, tb_log_dir, writer_dict)if perf_indicator > best_perf:best_perf = perf_indicatorbest_model = Trueelse:best_model = Falselogger.info('=> saving checkpoint to {}'.format(final_output_dir))save_checkpoint({'epoch': epoch + 1,'model': config.MODEL.NAME,'state_dict': model.state_dict(),'perf': perf_indicator,'optimizer': optimizer.state_dict(),}, best_model, final_output_dir, filename='checkpoint.pth.tar')final_model_state_file = os.path.join(final_output_dir,'final_state.pth.tar')logger.info('saving final model state to {}'.format(final_model_state_file))torch.save(model.state_dict(), final_model_state_file)#torch.dule.state_dict(), final_model_state_file)writer_dict['writer'].close()final_pth_file = os.path.join(final_output_dir,  'HRNet.pth')#增加的代码print("final_pth_file:", final_pth_file)#增加的代码torch.save(model, final_pth_file)#torch.dule, final_pth_file)#增加的代码if __name__ == '__main__':main()

/tools/valid.py

# ------------------------------------------------------------------------------
# Copyright (c) Microsoft
# Licensed under the MIT License.
# Written by Bin Xiao (Bin.Xiao@microsoft)
# Modified by Ke Sun (sunk@mail.ustc.edu)
# ------------------------------------------------------------------------------from __future__ import absolute_import
from __future__ import division
from __future__ import print_functionimport argparse
import os
#os.environ['CUDA_VISIBLE_DEVICES']='1/'
import sys
import shutil
import pprintimport torch
parallel
import torch.backends.cudnn as cudnn
import torch.optim
import torch.utils.data
import torch.utils.data.distributed
import torchvision.datasets as datasets
ansforms as transformsimport _init_paths
import models
from config import config
from config import update_config
from core.function import validate
delsummary import get_model_summary
from utils.utils import create_loggerdef parse_args():parser = argparse.ArgumentParser(description='Train keypoints network')parser.add_argument('--cfg',help='experiment configure file name',required=True,type=str)parser.add_argument('--modelDir',help='model directory',type=str,default='')parser.add_argument('--logDir',help='log directory',type=str,default='')parser.add_argument('--dataDir',help='data directory',type=str,default='')parser.add_argument('--testModel',help='testModel',type=str,default='')args = parser.parse_args()update_config(config, args)return argsdef main():args = parse_args()logger, final_output_dir, tb_log_dir = create_logger(config, args.cfg, 'valid')logger.info(pprint.pformat(args))logger.info(pprint.pformat(config))# cudnn related settingcudnn.benchmark = config.CUDNN.BENCHMARKtorch.backends.cudnn.deterministic = config.CUDNN.DETERMINISTICtorch.abled = config.CUDNN.ENABLEDmodel = eval('models.'+config.MODEL.NAME+'.get_cls_net')(config)dump_input = torch.rand((1, 3, config.MODEL.IMAGE_SIZE[1], config.MODEL.IMAGE_SIZE[0]))logger.info(get_model_summary(model, dump_input))if config.TEST.MODEL_FILE:logger.info('=> loading model from {}'.format(config.TEST.MODEL_FILE))_checkpoint=torch.load(config.TEST.MODEL_FILE);model.load_state_dict(_checkpoint)# model.load_state_dict(torch.load(config.TEST.MODEL_FILE))else:model_state_file = os.path.join(final_output_dir,'final_state.pth.tar')logger.info('=> loading model from {}'.format(model_state_file))model.load_state_dict(torch.load(model_state_file))gpus = list(config.GPUS)#model = DataParallel(model, device_ids=gpus).cuda()device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")model&#(device)# define loss function (criterion) and optimizercriterion = CrossEntropyLoss().cuda()# Data loading codevaldir = os.path.join(config.DATASET.ROOT,config.DATASET.TEST_SET)normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],std=[0.229, 0.224, 0.225])valid_loader = torch.utils.data.DataLoader(datasets.ImageFolder(valdir, transforms.Compose([transforms.Resize(int(config.MODEL.IMAGE_SIZE[0] / 0.875)),transforms.CenterCrop(config.MODEL.IMAGE_SIZE[0]),transforms.ToTensor(),normalize,])),batch_size=config.TEST.BATCH_SIZE_PER_GPU*len(gpus),shuffle=False,num_workers=config.WORKERS,pin_memory=True)# evaluate on validation setvalidate(config, valid_loader, model, criterion, final_output_dir,tb_log_dir, None)if __name__ == '__main__':main()

/lib/config/default.py

# ------------------------------------------------------------------------------
# Copyright (c) Microsoft
# Licensed under the MIT License.
# Written by Bin Xiao (Bin.Xiao@microsoft)
# Modified by Ke Sun (sunk@mail.ustc.edu)
# ------------------------------------------------------------------------------from __future__ import absolute_import
from __future__ import division
from __future__ import print_functionimport fig import CfgNode as CN_C = CN()_C.OUTPUT_DIR = ''
_C.LOG_DIR = ''
_C.DATA_DIR = ''
_C.GPUS = (0,)
_C.WORKERS = 4
_C.PRINT_FREQ = 20
_C.AUTO_RESUME = False
_C.PIN_MEMORY = True
_C.RANK = 0# Cudnn related params
_C.CUDNN = CN()
_C.CUDNN.BENCHMARK = True
_C.CUDNN.DETERMINISTIC = False
_C.CUDNN.ENABLED = True# common params for NETWORK
_C.MODEL = CN()
_C.MODEL.NAME = 'cls_hrnet'
_C.MODEL.INIT_WEIGHTS = True
_C.MODEL.PRETRAINED = ''
_C.MODEL.NUM_JOINTS = 17
_C.MODEL.NUM_CLASSES = 1000
_C.MODEL.TAG_PER_JOINT = True
_C.MODEL.TARGET_TYPE = 'gaussian'
_C.MODEL.IMAGE_SIZE = [256, 256]  # width * height, ex: 192 * 256
_C.MODEL.HEATMAP_SIZE = [64, 64]  # width * height, ex: 24 * 32
_C.MODEL.SIGMA = 2
_C.MODEL.EXTRA = CN(new_allowed=True)_C.LOSS = CN()
_C.LOSS.USE_OHKM = False
_C.LOSS.TOPK = 8
_C.LOSS.USE_TARGET_WEIGHT = True
_C.LOSS.USE_DIFFERENT_JOINTS_WEIGHT = False# DATASET related params
_C.DATASET = CN()
_C.DATASET.ROOT = ''
_C.DATASET.DATASET = 'mpii'
_C.DATASET.TRAIN_SET = 'train'
_C.DATASET.TEST_SET = 'valid'
_C.DATASET.DATA_FORMAT = 'jpg'
_C.DATASET.HYBRID_JOINTS_TYPE = ''
_C.DATASET.SELECT_DATA = False# training data augmentation
_C.DATASET.FLIP = True
_C.DATASET.SCALE_FACTOR = 0.25
_C.DATASET.ROT_FACTOR = 30
_C.DATASET.PROB_HALF_BODY = 0.0
_C.DATASET.NUM_JOINTS_HALF_BODY = 8
_C.DATASET.COLOR_RGB = False# train
_C.TRAIN = CN()_C.TRAIN.LR_FACTOR = 0.1
_C.TRAIN.LR_STEP = [90, 110]
_C.TRAIN.LR = 0.001_C.TRAIN.OPTIMIZER = 'adam'
_C.TRAIN.MOMENTUM = 0.9
_C.TRAIN.WD = 0.0001
_C.TRAIN.NESTEROV = False
_C.TRAIN.GAMMA1 = 0.99
_C.TRAIN.GAMMA2 = 0.0_C.TRAIN.BEGIN_EPOCH = 0
_C.TRAIN.END_EPOCH =2# 140_C.TRAIN.RESUME = False
_C.TRAIN.CHECKPOINT = ''_C.TRAIN.BATCH_SIZE_PER_GPU = 32
_C.TRAIN.SHUFFLE = True# testing
_C.TEST = CN()# size of images for each device
_C.TEST.BATCH_SIZE_PER_GPU = 32
# Test Model Epoch
_C.TEST.FLIP_TEST = False
_C.TEST.POST_PROCESS = False
_C.TEST.SHIFT_HEATMAP = False_C.TEST.USE_GT_BBOX = False# nms
_C.TEST.IMAGE_THRE = 0.1
_C.TEST.NMS_THRE = 0.6
_C.TEST.SOFT_NMS = False
_C.TEST.OKS_THRE = 0.5
_C.TEST.IN_VIS_THRE = 0.0
_C.TEST.COCO_BBOX_FILE = ''
_C.TEST.BBOX_THRE = 1.0
_C.TEST.MODEL_FILE = ''# debug
_C.DEBUG = CN()
_C.DEBUG.DEBUG = False
_C.DEBUG.SAVE_BATCH_IMAGES_GT = False
_C.DEBUG.SAVE_BATCH_IMAGES_PRED = False
_C.DEBUG.SAVE_HEATMAPS_GT = False
_C.DEBUG.SAVE_HEATMAPS_PRED = Falsedef update_config(cfg, args):cfg.defrost()_from_file(args.cfg)delDir:cfg.OUTPUT_DIR = delDirif args.logDir:cfg.LOG_DIR = args.logDirif args.dataDir:cfg.DATA_DIR = args.stModel:cfg.TEST.MODEL_FILE = stModelcfg.DATASET.ROOT = os.path.join(cfg.DATA_DIR, cfg.DATASET.DATASET, 'images')cfg.freeze()if __name__ == '__main__':import syswith open(sys.argv[1], 'w') as f:print(_C, file=f)

/experiments/cls_hrnet_w18_sgd_lr5e-2_wd1e-4_bs32_x100.yaml

GPUS: (0,1,2,3)
LOG_DIR: 'log/'
DATA_DIR: ''
OUTPUT_DIR: 'output/'
WORKERS: 4
PRINT_FREQ: 20MODEL: NAME: cls_hrnet IMAGE_SIZE: - 224- 224 EXTRA: STAGE1:NUM_MODULES: 1NUM_RANCHES: 1BLOCK: BOTTLENECKNUM_BLOCKS:- 4NUM_CHANNELS:- 64FUSE_METHOD: SUMSTAGE2: NUM_MODULES: 1NUM_BRANCHES: 2BLOCK: BASICNUM_BLOCKS- 4- 4NUM_CHANNELS:- 18- 36FUSE_METHOD: SUMSTAGE3:NUM_MODULES: 4NUM_BRANCHES: 3BLOCK: BASICNUM_BLOCKS:- 4- 4- 4NUM_CHANNELS:- 18- 36- 72FUSE_METHOD: SUMSTAGE4:NUM_MODULES: 3NUM_BRANCHES: 4BLOCK: BASICNUM_BLOCKS:- 4- 4- 4- 4NUM_CHANNELS:- 18- 36- 72- 144FUSE_METHOD: SUM
CUDNN:BENCHMARK: trueDETERMINISTIC: falseENABLED: true
DATASET:DATASET: 'imagenet'DATA_FORMAT: 'jpg'ROOT: 'data/imagenet/'TEST_SET: 'val'TRAIN_SET: 'train'
TEST:BATCH_SIZE_PER_GPU: 32MODEL_FILE: ''
TRAIN:BATCH_SIZE_PER_GPU: 32BEGIN_EPOCH: 0END_EPOCH: 2RESUME: trueLR_FACTOR: 0.1LR_FACTOR: 0.1LR_STEP:- 30- 60- 90OPTIMIZER: sgdLR: 0.05WD: 0.0001MOMENTUM: 0.9NESTEROV: trueSHUFFLE: true
DEBUG:DEBUG: false