def provide_model(model_name, classes_count, device, full=True):
"""
provides the pretrained model by model_name and replaces last classification layer with number of classes
if full is provided and set to False, then the weights of the pretrained are set unmodifiable
"""
if model_name == 'hrnet':
hrnet_model = load_hrnet_pretrained(device)
num_ftrs = hrnet_model.classifier.in_features
hrnet_model.classifier = nn.Linear(num_ftrs, classes_count)
model = nn.Sequential(hrnet_model).to(device)
elif model_name == 'resnet18':
resnet18 = torchvision.models.resnet18(pretrained=True)
num_ftrs = resnet18.fc.in_features
resnet18.fc = nn.Linear(num_ftrs, classes_count)
model = resnet18
elif model_name == 'resnet152':
resnet152 = torchvision.models.resnet152(pretrained=True)
num_ftrs = resnet152.fc.in_features
resnet152.fc = nn.Linear(num_ftrs, classes_count)
model = resnet152
dump_input = torch.rand(
(1, 3, config.MODEL.IMAGE_SIZE[1], config.MODEL.IMAGE_SIZE[0]))
print(get_model_summary(model.cpu(), dump_input.cpu(), verbose=True))
return model
def 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 setting
cudnn.benchmark = config.CUDNN.BENCHMARK
torch.backends.cudnn.deterministic = config.CUDNN.DETERMINISTIC
torch.backends.cudnn.enabled = config.CUDNN.ENABLED
model = 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))
# model.load_state_dict(torch.load(config.TEST.MODEL_FILE))
model.load_state_dict({
k.replace('module.', ''): v
for k, v in torch.load(config.TEST.MODEL_FILE)
['state_dict'].items()
})
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 = [0, 1]
model = torch.nn.DataParallel(model, device_ids=gpus).cuda()
# define loss function (criterion) and optimizer
criterion = torch.nn.CrossEntropyLoss().cuda()
# Data loading code
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])
test_dataset = datasets(dataset_root='./Data/',
split='test',
size=config.MODEL.IMAGE_SIZE[0])
valid_loader = torch.utils.data.DataLoader(
test_dataset,
batch_size=config.TRAIN.BATCH_SIZE_PER_GPU * len(gpus),
shuffle=True,
num_workers=config.WORKERS,
pin_memory=True)
# evaluate on validation set
validate(config, valid_loader, model, criterion, final_output_dir,
tb_log_dir, None)
def main():
mean = [0.485, 0.456, 0.406],
std = [0.229, 0.224, 0.225]
args = parse_args()
logger, final_output_dir, _ = create_logger(config, args.cfg, 'test')
logger.info(pprint.pformat(args))
logger.info(pprint.pformat(config))
# cudnn related setting
cudnn.benchmark = config.CUDNN.BENCHMARK
cudnn.deterministic = config.CUDNN.DETERMINISTIC
cudnn.enabled = config.CUDNN.ENABLED
# build model
if torch.__version__.startswith('1'):
module = eval('models.' + config.MODEL.NAME)
module.BatchNorm2d_class = module.BatchNorm2d = torch.nn.BatchNorm2d
model = eval('models.' + config.MODEL.NAME + '.get_seg_model')(config)
dump_input = torch.rand(
(1, 3, config.TRAIN.IMAGE_SIZE[1], config.TRAIN.IMAGE_SIZE[0]))
logger.info(get_model_summary(model.cuda(), dump_input.cuda()))
if config.TEST.MODEL_FILE:
model_state_file = config.TEST.MODEL_FILE
else:
# model_state_file = os.path.join(final_output_dir, 'best_0.7589.pth')
model_state_file = os.path.join(final_output_dir, 'best.pth')
logger.info('=> loading model from {}'.format(model_state_file))
pretrained_dict = torch.load(
'/home/hwits/Documents/CarVid/DDRNet/DDRNet.pytorch/model_best_bacc.pth.tar'
)
if 'state_dict' in pretrained_dict:
pretrained_dict = pretrained_dict['state_dict']
newstate_dict = {
k: v
for k, v in pretrained_dict.items() if k in model.state_dict()
}
# print(pretrained_dict.keys())
model.load_state_dict(newstate_dict)
model = model.cuda()
if True:
save_wts = True
print(model)
if save_wts:
f = open('DDRNetLite.wts', 'w')
f.write('{}\n'.format(len(model.state_dict().keys())))
for k, v in model.state_dict().items():
print("Layer {} ; Size {}".format(k, v.cpu().numpy().shape))
vr = v.reshape(-1).cpu().numpy()
f.write('{} {} '.format(k, len(vr)))
for vv in vr:
f.write(' ')
f.write(struct.pack('>f', float(vv)).hex())
f.write('\n')
def validate_hrnet():
# cudnn related setting
torch.backends.cudnn.benchmark = config.CUDNN.BENCHMARK
torch.backends.cudnn.deterministic = config.CUDNN.DETERMINISTIC
torch.backends.cudnn.enabled = config.CUDNN.ENABLED
model = eval('models.' + config.MODEL.NAME + '.get_cls_net')(config)
device = torch.device("cuda:0")
dump_input = torch.rand(
(1, 3, config.MODEL.IMAGE_SIZE[1], config.MODEL.IMAGE_SIZE[0]))
print(
get_model_summary(model.cuda(device),
dump_input.cuda(device),
verbose=True))
if config.TEST.MODEL_FILE:
print('=> loading model from {}'.format(config.TEST.MODEL_FILE))
model.load_state_dict(torch.load(config.TEST.MODEL_FILE))
else:
model_state_file = os.path.join("./", 'final_state.pth.tar')
print('=> loading model from {}'.format(model_state_file))
model.load_state_dict(torch.load(model_state_file))
# model = torch.nn.Sequential(model).cuda()
model = torch.nn.Sequential(model).cuda(device)
# define loss function (criterion) and optimizer
criterion = torch.nn.CrossEntropyLoss().cuda(device)
# Data loading code
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])
dataset = 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,
]))
print(f'number of images: {len(dataset.imgs)}')
valid_loader = torch.utils.data.DataLoader(
dataset,
batch_size=config.TEST.BATCH_SIZE_PER_GPU,
shuffle=False,
num_workers=config.WORKERS,
pin_memory=True)
# evaluate on validation set
validate(config, valid_loader, model, criterion, './log', '/output', None)
def main():
mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]
args = parse_args()
logger, final_output_dir, _ = create_logger(
config, args.cfg, 'test')
logger.info(pprint.pformat(args))
logger.info(pprint.pformat(config))
# cudnn related setting
cudnn.benchmark = config.CUDNN.BENCHMARK
cudnn.deterministic = config.CUDNN.DETERMINISTIC
cudnn.enabled = config.CUDNN.ENABLED
# build model
if torch.__version__.startswith('1'):
module = eval('models.'+config.MODEL.NAME)
module.BatchNorm2d_class = module.BatchNorm2d = torch.nn.BatchNorm2d
model = eval('models.'+config.MODEL.NAME +
'.get_seg_model')(config)
dump_input = torch.rand(
(1, 3, config.TRAIN.IMAGE_SIZE[1], config.TRAIN.IMAGE_SIZE[0])
)
logger.info(get_model_summary(model.cuda(), dump_input.cuda()))
if config.TEST.MODEL_FILE:
model_state_file = config.TEST.MODEL_FILE
else:
# model_state_file = os.path.join(final_output_dir, 'best_0.7589.pth')
model_state_file = os.path.join(final_output_dir, 'best.pth')
logger.info('=> loading model from {}'.format(model_state_file))
pretrained_dict = torch.load('/home/kong/Documents/DDRNet.Pytorch/DDRNet.Pytorch/output/face/ddrnet23_slim/checkpoint.pth.tar')
if 'state_dict' in pretrained_dict:
pretrained_dict = pretrained_dict['state_dict']
newstate_dict = {k:v for k,v in pretrained_dict.items() if k in model.state_dict()}
# print(pretrained_dict.keys())
model.load_state_dict(newstate_dict)
model = model.to("cpu")
print(model)
model.eval()
example = torch.rand(1, 3, 512, 512)
model = torch.quantization.convert(model)
# traced_script_module = torch.jit.trace(model, example)
# traced_script_module.save("ddrnetfp32.pt")
scriptedm = torch.jit.script(model)
opt_model = torch.utils.optimize_for_mobile(scriptedm)
torch.jit.save(opt_model, "ddrnetint8.pt")
def load_hrnet_pretrained(device):
# cudnn related setting
torch.backends.cudnn.benchmark = config.CUDNN.BENCHMARK
torch.backends.cudnn.deterministic = config.CUDNN.DETERMINISTIC
torch.backends.cudnn.enabled = config.CUDNN.ENABLED
model = 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]))
if device.type == 'cpu':
print(get_model_summary(model, dump_input, verbose=True))
else:
print(
get_model_summary(model.cuda(device),
dump_input.cuda(device),
verbose=True))
print('=> loading model from {}'.format(config.TEST.MODEL_FILE))
model.load_state_dict(torch.load(config.TEST.MODEL_FILE))
return model
def main():
args = parse_args()
if args.dist_url == "env://" and args.world_size == -1:
args.world_size = int(os.environ["WORLD_SIZE"])
dist.init_process_group(backend=args.dist_backend,
init_method=args.dist_url)
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 setting
cudnn.benchmark = config.CUDNN.BENCHMARK
torch.backends.cudnn.deterministic = config.CUDNN.DETERMINISTIC
torch.backends.cudnn.enabled = config.CUDNN.ENABLED
model = 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 file
# this_dir = os.path.dirname(__file__)
# models_dst_dir = os.path.join(final_output_dir, 'models')
# if os.path.exists(models_dst_dir):
# shutil.rmtree(models_dst_dir)
# shutil.copytree(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 = torch.nn.DataParallel(model, device_ids=gpus).cuda()
'''
# Change DP to DDP
torch.cuda.set_device(args.local_rank)
model = model.to(args.local_rank)
model = torch.nn.parallel.DistributedDataParallel(
model, device_ids=[args.local_rank], output_device=args.local_rank)
# define loss function (criterion) and optimizer
criterion = torch.nn.CrossEntropyLoss().cuda()
optimizer = get_optimizer(config, model)
best_perf = 0.0
best_model = False
last_epoch = config.TRAIN.BEGIN_EPOCH
if 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.module.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
logger.info("=> loaded checkpoint (epoch {})".format(
checkpoint['epoch']))
best_model = True
if 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 code
traindir = 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,
])
)
'''
# Change to TSV dataset instance
train_dataset = TSVInstance(
traindir,
transforms.Compose([
transforms.RandomResizedCrop(config.MODEL.IMAGE_SIZE[0]),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
]))
# DDP requires DistributedSampler
train_sampler = torch.utils.data.distributed.DistributedSampler(
train_dataset)
train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=config.TRAIN.BATCH_SIZE_PER_GPU,
shuffle=(train_sampler is None),
num_workers=config.WORKERS,
pin_memory=True,
sampler=train_sampler)
valid_loader = torch.utils.data.DataLoader(
TSVInstance(
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,
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 epoch
train(config, train_loader, model, criterion, optimizer, epoch,
final_output_dir, tb_log_dir, writer_dict)
# evaluate on validation set
perf_indicator = validate(config, valid_loader, model, criterion,
final_output_dir, tb_log_dir, writer_dict)
if perf_indicator > best_perf:
best_perf = perf_indicator
best_model = True
else:
best_model = False
logger.info('=> saving checkpoint to {}'.format(final_output_dir))
save_checkpoint(
{
'epoch': epoch + 1,
'model': config.MODEL.NAME,
'state_dict': model.module.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.module.state_dict(), final_model_state_file)
writer_dict['writer'].close()
def main():
args = parse_args()
print(colored("Setting default tensor type to cuda.FloatTensor", "cyan"))
torch.multiprocessing.set_start_method('spawn')
torch.set_default_tensor_type('torch.cuda.FloatTensor')
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 setting
cudnn.benchmark = config.CUDNN.BENCHMARK
torch.backends.cudnn.deterministic = config.CUDNN.DETERMINISTIC
torch.backends.cudnn.enabled = config.CUDNN.ENABLED
model = eval('models.' + config.MODEL.NAME + '.get_cls_net')(config).cuda()
dump_input = torch.rand(config.TRAIN.BATCH_SIZE_PER_GPU, 3,
config.MODEL.IMAGE_SIZE[1],
config.MODEL.IMAGE_SIZE[0]).cuda()
logger.info(get_model_summary(model, dump_input))
if config.TRAIN.MODEL_FILE:
model.load_state_dict(torch.load(config.TRAIN.MODEL_FILE))
logger.info(
colored('=> loading model from {}'.format(config.TRAIN.MODEL_FILE),
'red'))
# copy model file
this_dir = os.path.dirname(__file__)
models_dst_dir = os.path.join(final_output_dir, 'models')
if os.path.exists(models_dst_dir):
shutil.rmtree(models_dst_dir)
shutil.copytree(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 = nn.DataParallel(model, device_ids=gpus).cuda()
print("Finished constructing model!")
# define loss function (criterion) and optimizer
criterion = nn.CrossEntropyLoss().cuda()
optimizer = get_optimizer(config, model)
lr_scheduler = None
best_perf = 0.0
best_model = False
last_epoch = config.TRAIN.BEGIN_EPOCH
if 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.module.load_state_dict(checkpoint['state_dict'])
# Update weight decay if needed
checkpoint['optimizer']['param_groups'][0][
'weight_decay'] = config.TRAIN.WD
optimizer.load_state_dict(checkpoint['optimizer'])
if 'lr_scheduler' in checkpoint:
lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
optimizer,
1e5,
last_epoch=checkpoint['lr_scheduler']['last_epoch'])
lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
logger.info("=> loaded checkpoint (epoch {})".format(
checkpoint['epoch']))
best_model = True
# Data loading code
dataset_name = config.DATASET.DATASET
if dataset_name == 'imagenet':
traindir = os.path.join(config.DATASET.ROOT + '/images',
config.DATASET.TRAIN_SET)
valdir = os.path.join(config.DATASET.ROOT + '/images',
config.DATASET.TEST_SET)
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
transform_train = transforms.Compose([
transforms.RandomResizedCrop(config.MODEL.IMAGE_SIZE[0]),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
])
transform_valid = transforms.Compose([
transforms.Resize(int(config.MODEL.IMAGE_SIZE[0] / 0.875)),
transforms.CenterCrop(config.MODEL.IMAGE_SIZE[0]),
transforms.ToTensor(),
normalize,
])
train_dataset = datasets.ImageFolder(traindir, transform_train)
valid_dataset = datasets.ImageFolder(valdir, transform_valid)
else:
assert dataset_name == "cifar10", "Only CIFAR-10 is supported at this phase"
classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog',
'horse', 'ship', 'truck') # For reference
normalize = transforms.Normalize((0.4914, 0.4822, 0.4465),
(0.2023, 0.1994, 0.2010))
transform_train = transforms.Compose([
transforms.RandomCrop(
32,
padding=4), # ** Remove this line if no data augmentation **
transforms.RandomHorizontalFlip(
), # ** Remove this line if no data augmentation **
transforms.ToTensor(),
normalize,
])
transform_valid = transforms.Compose([
transforms.ToTensor(),
normalize,
])
train_dataset = datasets.CIFAR10(root=f'{config.DATASET.ROOT}',
train=True,
download=True,
transform=transform_train)
valid_dataset = datasets.CIFAR10(root=f'{config.DATASET.ROOT}',
train=False,
download=True,
transform=transform_valid)
train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=config.TRAIN.BATCH_SIZE_PER_GPU * len(gpus),
shuffle=True,
num_workers=config.WORKERS,
pin_memory=True)
valid_loader = torch.utils.data.DataLoader(
valid_dataset,
batch_size=config.TEST.BATCH_SIZE_PER_GPU * len(gpus),
shuffle=False,
num_workers=config.WORKERS,
pin_memory=True)
# Learning rate scheduler
if lr_scheduler is None:
if config.TRAIN.LR_SCHEDULER != 'step':
lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
optimizer,
len(train_loader) * config.TRAIN.END_EPOCH,
eta_min=1e-6)
elif 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)
# Training code
for epoch in range(last_epoch, config.TRAIN.END_EPOCH):
topk = (1, 5) if dataset_name == 'imagenet' else (1, )
if config.TRAIN.LR_SCHEDULER == 'step':
lr_scheduler.step()
# train for one epoch
train(config,
train_loader,
model,
criterion,
optimizer,
lr_scheduler,
epoch,
final_output_dir,
tb_log_dir,
writer_dict,
topk=topk)
torch.cuda.empty_cache()
# evaluate on validation set
perf_indicator = validate(config,
valid_loader,
model,
criterion,
lr_scheduler,
epoch,
final_output_dir,
tb_log_dir,
writer_dict,
topk=topk)
torch.cuda.empty_cache()
writer_dict['writer'].flush()
if perf_indicator > best_perf:
best_perf = perf_indicator
best_model = True
else:
best_model = False
logger.info('=> saving checkpoint to {}'.format(final_output_dir))
save_checkpoint(
{
'epoch': epoch + 1,
'model': config.MODEL.NAME,
'state_dict': model.module.state_dict(),
'perf': perf_indicator,
'optimizer': optimizer.state_dict(),
'lr_scheduler': lr_scheduler.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.module.state_dict(), final_model_state_file)
writer_dict['writer'].close()
def 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 setting
cudnn.benchmark = config.CUDNN.BENCHMARK
torch.backends.cudnn.deterministic = config.CUDNN.DETERMINISTIC
torch.backends.cudnn.enabled = config.CUDNN.ENABLED
model = 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 file
this_dir = os.path.dirname(__file__)
models_dst_dir = os.path.join(final_output_dir, 'models')
if os.path.exists(models_dst_dir):
shutil.rmtree(models_dst_dir)
shutil.copytree(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 = torch.nn.DataParallel(model, device_ids=gpus).cuda()
# define loss function (criterion) and optimizer
criterion = torch.nn.CrossEntropyLoss().cuda()
optimizer = get_optimizer(config, model)
best_perf = 0.0
best_model = False
last_epoch = config.TRAIN.BEGIN_EPOCH
if 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.module.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
logger.info("=> loaded checkpoint (epoch {})".format(
checkpoint['epoch']))
best_model = True
if 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 code
traindir = 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("class.txt", "w") as f1:
for classname in train_dataset.classes:
f1.write(classname + "\n")
#print(train_dataset.class_to_idx) #按顺序为这些类别定义索引为0,1...
with open("classToIndex.txt", "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=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=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 epoch
train(config, train_loader, model, criterion, optimizer, epoch,
final_output_dir, tb_log_dir, writer_dict)
# evaluate on validation set
perf_indicator = validate(config, valid_loader, model, criterion,
final_output_dir, tb_log_dir, writer_dict)
if perf_indicator > best_perf:
best_perf = perf_indicator
best_model = True
else:
best_model = False
logger.info('=> saving checkpoint to {}'.format(final_output_dir))
save_checkpoint(
{
'epoch': epoch + 1,
'model': config.MODEL.NAME,
'state_dict': model.module.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.module.state_dict(), final_model_state_file)
final_pth_file = os.path.join(final_output_dir, 'HRNet.pt')
print("final_pth_file:", final_pth_file)
torch.save(model.module, final_pth_file)
writer_dict['writer'].close()
def 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 setting
cudnn.benchmark = config.CUDNN.BENCHMARK
torch.backends.cudnn.deterministic = config.CUDNN.DETERMINISTIC
torch.backends.cudnn.enabled = config.CUDNN.ENABLED
model = eval('models.' + config.MODEL.NAME +
'.get_contrastive_net')(config).cuda()
dump_input = torch.rand(
(1, 3, config.MODEL.IMAGE_SIZE[1], config.MODEL.IMAGE_SIZE[0])).cuda()
logger.info(get_model_summary(model, dump_input))
if config.TEST.MODEL_FILE:
logger.info('=> loading model from {}'.format(config.TEST.MODEL_FILE))
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 = torch.nn.DataParallel(model, device_ids=gpus).cuda()
# define loss function (criterion) and optimizer
criterion = torch.nn.CrossEntropyLoss().cuda()
# Data loading code
valdir = os.path.join(config.DATASET.ROOT + '/images',
config.DATASET.TEST_SET)
dataset_name = config.DATASET.DATASET
if dataset_name == 'imagenet':
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)
else:
normalize = transforms.Normalize((0.4914, 0.4822, 0.4465),
(0.2023, 0.1994, 0.2010))
transform_valid = transforms.Compose([
transforms.ToTensor(),
normalize,
])
if dataset_name == 'cifar10':
valid_dataset = datasets.CIFAR10(root=f'{config.DATASET.ROOT}',
train=False,
download=True,
transform=transform_valid)
elif dataset_name == 'cifar100':
valid_dataset = datasets.CIFAR100(root=f'{config.DATASET.ROOT}',
train=False,
download=True,
transform=transform_valid)
valid_loader = torch.utils.data.DataLoader(
valid_dataset,
batch_size=config.TEST.BATCH_SIZE_PER_GPU * len(gpus),
shuffle=False,
num_workers=config.WORKERS,
pin_memory=True)
model.eval()
embs = []
targets = []
with torch.no_grad():
for i, (input, target) in enumerate(valid_loader):
output = model(input, train_step=-1)[-1].cpu().numpy()
embs.extend(output)
targets.extend(target.cpu().numpy())
embs = np.asarray(embs)
print(targets, flush=True)
targets = np.asarray(targets)
pca = PCA(n_components=64)
embs = pca.fit_transform(embs)
tsne = TSNE(n_components=2, random_state=0)
embs_2d = tsne.fit_transform(embs)
target_ids = range(10)
plt.figure(figsize=(10, 10))
colors = 'r', 'g', 'b', 'c', 'm', 'yellow', 'k', 'lime', 'orange', 'purple'
classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse',
'ship', 'truck') # For reference
for i, c, label in zip(target_ids, colors, classes):
plt.scatter(embs_2d[targets == i, 0],
embs_2d[targets == i, 1],
c=c,
label=label)
plt.legend()
plt.savefig(os.path.join(final_output_dir, 'tsne.png'))
def main():
args = parse_args()
logger, final_output_dir, tb_log_dir = create_logger(
config, args.cfg, 'train')
logger.info(pprint.pformat(args))
logger.info(config)
writer_dict = {
'writer': SummaryWriter(tb_log_dir),
'train_global_steps': 0,
'valid_global_steps': 0,
}
# cudnn related setting
cudnn.benchmark = config.CUDNN.BENCHMARK
cudnn.deterministic = config.CUDNN.DETERMINISTIC
cudnn.enabled = config.CUDNN.ENABLED
gpus = list(config.GPUS)
distributed = len(gpus) > 1
device = torch.device('cuda:{}'.format(args.local_rank))
# build model
model = eval('models.' + config.MODEL.NAME + '.get_seg_model')(config)
if args.local_rank == 0:
logger.info(model)
tot_params = sum(p.numel() for p in model.parameters()) / 1000000.0
logger.info(f">>> total params: {tot_params:.2f}M")
# provide the summary of model
dump_input = torch.rand(
(1, 3, config.TRAIN.IMAGE_SIZE[0], config.TRAIN.IMAGE_SIZE[1]))
logger.info(get_model_summary(model.to(device), dump_input.to(device)))
# copy model file
this_dir = os.path.dirname(__file__)
models_dst_dir = os.path.join(final_output_dir, 'models')
if os.path.exists(models_dst_dir):
shutil.rmtree(models_dst_dir)
shutil.copytree(os.path.join(this_dir, '../lib/models'),
models_dst_dir)
if distributed:
torch.cuda.set_device(args.local_rank)
torch.distributed.init_process_group(
backend="nccl",
init_method="env://",
)
# prepare data
train_dataset = eval('datasets.' + config.DATASET.DATASET)(
root=config.DATASET.ROOT,
list_path=config.DATASET.TRAIN_SET,
num_samples=None,
num_classes=config.DATASET.NUM_CLASSES,
multi_scale=config.TRAIN.MULTI_SCALE,
flip=config.TRAIN.FLIP,
ignore_label=config.TRAIN.IGNORE_LABEL,
base_size=config.TRAIN.BASE_SIZE,
crop_size=tuple(config.TRAIN.IMAGE_SIZE), # (height, width)
scale_factor=config.TRAIN.SCALE_FACTOR)
if distributed:
train_sampler = DistributedSampler(train_dataset)
else:
train_sampler = None
trainloader = torch.utils.data.DataLoader(
train_dataset,
batch_size=config.TRAIN.BATCH_SIZE_PER_GPU,
shuffle=config.TRAIN.SHUFFLE and train_sampler is None,
num_workers=config.WORKERS,
pin_memory=True,
drop_last=True,
sampler=train_sampler)
test_dataset = eval('datasets.' + config.DATASET.DATASET)(
root=config.DATASET.ROOT,
list_path=config.DATASET.TEST_SET,
num_samples=config.TEST.NUM_SAMPLES,
num_classes=config.DATASET.NUM_CLASSES,
multi_scale=False,
flip=False,
ignore_label=config.TRAIN.IGNORE_LABEL,
base_size=config.TEST.BASE_SIZE,
crop_size=tuple(config.TEST.IMAGE_SIZE), # (height, width)
)
if distributed:
test_sampler = DistributedSampler(test_dataset)
else:
test_sampler = None
testloader = torch.utils.data.DataLoader(
test_dataset,
batch_size=config.TEST.BATCH_SIZE_PER_GPU,
shuffle=False,
num_workers=config.WORKERS,
pin_memory=True,
sampler=test_sampler)
# criterion
if config.LOSS.USE_OHEM:
criterion = OhemCrossEntropy(ignore_label=config.TRAIN.IGNORE_LABEL,
weight=train_dataset.class_weights,
thresh=config.LOSS.OHEMTHRESH,
min_kept=config.LOSS.OHEMKEEP)
else:
criterion = CrossEntropy(ignore_label=config.TRAIN.IGNORE_LABEL,
weight=train_dataset.class_weights)
model_state_file = config.MODEL.PRETRAINED
logger.info('=> Loading model from {}'.format(model_state_file))
pretrained_dict = torch.load(model_state_file)
model_dict = model.state_dict()
pretrained_dict = {
k[6:]: v
for k, v in pretrained_dict.items() if k[6:] in model_dict.keys()
}
for k, _ in pretrained_dict.items():
logger.info('=> Loading {} from pretrained model'.format(k))
model_dict.update(pretrained_dict)
model.load_state_dict(model_dict)
model = FullModel(model, criterion)
if distributed:
model = nn.SyncBatchNorm.convert_sync_batchnorm(model)
model = model.to(device)
if distributed:
model = nn.parallel.DistributedDataParallel(
model, device_ids=[args.local_rank], output_device=args.local_rank)
# optimizer
optimizer = get_optimizer(config, model)
epoch_iters = np.int(train_dataset.__len__() /
config.TRAIN.BATCH_SIZE_PER_GPU / len(gpus))
best_mIoU = 0
last_epoch = 0
if 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,
map_location=lambda storage, loc: storage)
best_mIoU = checkpoint['best_mIoU']
last_epoch = checkpoint['epoch']
model.module.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
logger.info("=> loaded checkpoint (epoch {})".format(
checkpoint['epoch']))
start = timeit.default_timer()
end_epoch = config.TRAIN.END_EPOCH
num_iters = config.TRAIN.END_EPOCH * epoch_iters
# learning rate scheduler
lr_scheduler_dict = {
'optimizer': optimizer,
'milestones': [s * epoch_iters for s in config.TRAIN.LR_STEP],
'gamma': config.TRAIN.LR_FACTOR,
'max_iters': num_iters,
'last_epoch': last_epoch,
'epoch_iters': epoch_iters
}
lr_scheduler = get_lr_scheduler(config.TRAIN.LR_SCHEDULER,
**lr_scheduler_dict)
for epoch in range(last_epoch, end_epoch):
if distributed:
train_sampler.set_epoch(epoch)
train(config, epoch, end_epoch, epoch_iters, trainloader, optimizer,
lr_scheduler, model, writer_dict, device)
valid_loss, mean_IoU = validate(config, testloader, model, writer_dict,
device)
if args.local_rank == 0:
logger.info(
'=> saving checkpoint to {}'.format(final_output_dir +
'/checkpoint.pth.tar'))
torch.save(
{
'epoch': epoch + 1,
'best_mIoU': best_mIoU,
'state_dict': model.module.state_dict(),
'optimizer': optimizer.state_dict()
}, os.path.join(final_output_dir, 'checkpoint.pth.tar'))
if mean_IoU > best_mIoU:
best_mIoU = mean_IoU
torch.save(model.module.state_dict(),
os.path.join(final_output_dir, 'best.pth'))
msg = f'Loss: {valid_loss:.4f}, MeanIU: {mean_IoU: 4.4f}, \
Best_mIoU: {best_mIoU: 4.4f}'
logger.info(msg)
if epoch == end_epoch - 1:
torch.save(model.module.state_dict(),
os.path.join(final_output_dir, 'final_state.pth'))
writer_dict['writer'].close()
end = timeit.default_timer()
logger.info(f'Hours: {np.int((end-start)/3600)}')
logger.info('Done!')
def main():
args = parse_args()
logger, final_output_dir, tb_log_dir = create_logger(
config, args.cfg, 'train')
logger.info(pprint.pformat(args))
logger.info(config)
writer_dict = {
'writer': SummaryWriter(tb_log_dir),
'train_global_steps': 0,
'valid_global_steps': 0,
}
# cudnn related setting
cudnn.benchmark = config.CUDNN.BENCHMARK
cudnn.deterministic = config.CUDNN.DETERMINISTIC
cudnn.enabled = config.CUDNN.ENABLED
gpus = list(config.GPUS)
distributed = len(gpus) > 1
device = torch.device('cuda:{}'.format(args.local_rank))
# build model
model = eval('models.' + config.MODEL.NAME + '.get_seg_model')(config)
if args.local_rank == 0:
# provide the summary of model
dump_input = torch.rand(
(1, 3, config.TRAIN.IMAGE_SIZE[1], config.TRAIN.IMAGE_SIZE[0]))
logger.info(get_model_summary(model.to(device), dump_input.to(device)))
# copy model file
this_dir = os.path.dirname(__file__)
models_dst_dir = os.path.join(final_output_dir, 'models')
if os.path.exists(models_dst_dir):
shutil.rmtree(models_dst_dir)
shutil.copytree(os.path.join(this_dir, '../lib/models'),
models_dst_dir)
if distributed:
torch.cuda.set_device(args.local_rank)
torch.distributed.init_process_group(
backend="nccl",
init_method="env://",
)
# prepare data
crop_size = (config.TRAIN.IMAGE_SIZE[1], config.TRAIN.IMAGE_SIZE[0])
train_dataset = eval('datasets.' + config.DATASET.DATASET)(
root=config.DATASET.ROOT,
list_path=config.DATASET.TRAIN_SET,
num_samples=None,
num_classes=config.DATASET.NUM_CLASSES,
multi_scale=config.TRAIN.MULTI_SCALE,
flip=config.TRAIN.FLIP,
ignore_label=config.TRAIN.IGNORE_LABEL,
base_size=config.TRAIN.BASE_SIZE,
crop_size=crop_size,
downsample_rate=config.TRAIN.DOWNSAMPLERATE,
scale_factor=config.TRAIN.SCALE_FACTOR)
if distributed:
train_sampler = DistributedSampler(train_dataset)
else:
train_sampler = None
trainloader = torch.utils.data.DataLoader(
train_dataset,
batch_size=config.TRAIN.BATCH_SIZE_PER_GPU,
shuffle=config.TRAIN.SHUFFLE and train_sampler is None,
num_workers=config.WORKERS,
pin_memory=True,
drop_last=True,
sampler=train_sampler)
if config.DATASET.EXTRA_TRAIN_SET:
extra_train_dataset = eval('datasets.' + config.DATASET.DATASET)(
root=config.DATASET.ROOT,
list_path=config.DATASET.EXTRA_TRAIN_SET,
num_samples=None,
num_classes=config.DATASET.NUM_CLASSES,
multi_scale=config.TRAIN.MULTI_SCALE,
flip=config.TRAIN.FLIP,
ignore_label=config.TRAIN.IGNORE_LABEL,
base_size=config.TRAIN.BASE_SIZE,
crop_size=crop_size,
downsample_rate=config.TRAIN.DOWNSAMPLERATE,
scale_factor=config.TRAIN.SCALE_FACTOR)
if distributed:
extra_train_sampler = DistributedSampler(extra_train_dataset)
else:
extra_train_sampler = None
extra_trainloader = torch.utils.data.DataLoader(
extra_train_dataset,
batch_size=config.TRAIN.BATCH_SIZE_PER_GPU,
shuffle=config.TRAIN.SHUFFLE and extra_train_sampler is None,
num_workers=config.WORKERS,
pin_memory=True,
drop_last=True,
sampler=extra_train_sampler)
test_size = (config.TEST.IMAGE_SIZE[1], config.TEST.IMAGE_SIZE[0])
test_dataset = eval('datasets.' + config.DATASET.DATASET)(
root=config.DATASET.ROOT,
list_path=config.DATASET.TEST_SET,
num_samples=config.TEST.NUM_SAMPLES,
num_classes=config.DATASET.NUM_CLASSES,
multi_scale=False,
flip=False,
ignore_label=config.TRAIN.IGNORE_LABEL,
base_size=config.TEST.BASE_SIZE,
crop_size=test_size,
center_crop_test=config.TEST.CENTER_CROP_TEST,
downsample_rate=1)
if distributed:
test_sampler = DistributedSampler(test_dataset)
else:
test_sampler = None
testloader = torch.utils.data.DataLoader(
test_dataset,
batch_size=config.TEST.BATCH_SIZE_PER_GPU,
shuffle=False,
num_workers=config.WORKERS,
pin_memory=True,
sampler=test_sampler)
# criterion
if config.LOSS.USE_OHEM:
criterion = OhemCrossEntropy(ignore_label=config.TRAIN.IGNORE_LABEL,
thres=config.LOSS.OHEMTHRES,
min_kept=config.LOSS.OHEMKEEP,
weight=train_dataset.class_weights)
else:
criterion = CrossEntropy(ignore_label=config.TRAIN.IGNORE_LABEL,
weight=train_dataset.class_weights)
model = FullModel(model, criterion)
model = nn.SyncBatchNorm.convert_sync_batchnorm(model)
model = model.to(device)
model = nn.parallel.DistributedDataParallel(model,
device_ids=[args.local_rank],
output_device=args.local_rank)
# optimizer
if config.TRAIN.OPTIMIZER == 'sgd':
optimizer = torch.optim.SGD(
[{
'params': filter(lambda p: p.requires_grad,
model.parameters()),
'lr': config.TRAIN.LR
}],
lr=config.TRAIN.LR,
momentum=config.TRAIN.MOMENTUM,
weight_decay=config.TRAIN.WD,
nesterov=config.TRAIN.NESTEROV,
)
else:
raise ValueError('Only Support SGD optimizer')
epoch_iters = np.int(train_dataset.__len__() /
config.TRAIN.BATCH_SIZE_PER_GPU / len(gpus))
best_mIoU = 0
last_epoch = 0
if 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,
map_location=lambda storage, loc: storage)
best_mIoU = checkpoint['best_mIoU']
last_epoch = checkpoint['epoch']
model.module.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
logger.info("=> loaded checkpoint (epoch {})".format(
checkpoint['epoch']))
start = timeit.default_timer()
end_epoch = config.TRAIN.END_EPOCH + config.TRAIN.EXTRA_EPOCH
num_iters = config.TRAIN.END_EPOCH * epoch_iters
extra_iters = config.TRAIN.EXTRA_EPOCH * epoch_iters
for epoch in range(last_epoch, end_epoch):
if distributed:
train_sampler.set_epoch(epoch)
if epoch >= config.TRAIN.END_EPOCH:
train(config, epoch - config.TRAIN.END_EPOCH,
config.TRAIN.EXTRA_EPOCH, epoch_iters, config.TRAIN.EXTRA_LR,
extra_iters, extra_trainloader, optimizer, model,
writer_dict, device)
else:
train(config, epoch, config.TRAIN.END_EPOCH, epoch_iters,
config.TRAIN.LR, num_iters, trainloader, optimizer, model,
writer_dict, device)
valid_loss, mean_IoU, IoU_array = validate(config, testloader, model,
writer_dict, device)
if args.local_rank == 0:
logger.info(
'=> saving checkpoint to {}'.format(final_output_dir +
'checkpoint.pth.tar'))
torch.save(
{
'epoch': epoch + 1,
'best_mIoU': best_mIoU,
'state_dict': model.module.state_dict(),
'optimizer': optimizer.state_dict(),
}, os.path.join(final_output_dir, 'checkpoint.pth.tar'))
if mean_IoU > best_mIoU:
best_mIoU = mean_IoU
torch.save(model.module.state_dict(),
os.path.join(final_output_dir, 'best.pth'))
msg = 'Loss: {:.3f}, MeanIU: {: 4.4f}, Best_mIoU: {: 4.4f}'.format(
valid_loss, mean_IoU, best_mIoU)
logging.info(msg)
logging.info(IoU_array)
if epoch == end_epoch - 1:
torch.save(model.module.state_dict(),
os.path.join(final_output_dir, 'final_state.pth'))
writer_dict['writer'].close()
end = timeit.default_timer()
logger.info('Hours: %d' % np.int((end - start) / 3600))
logger.info('Done')
def main():
args = parse_args()
logger, final_output_dir, _ = create_logger(config, args.cfg, 'test')
logger.info(pprint.pformat(args))
logger.info(pprint.pformat(config))
# cudnn related setting
cudnn.benchmark = config.CUDNN.BENCHMARK
cudnn.deterministic = config.CUDNN.DETERMINISTIC
cudnn.enabled = config.CUDNN.ENABLED
# build model
if torch.__version__.startswith('1'):
module = eval('models.' + config.MODEL.NAME)
module.BatchNorm2d_class = module.BatchNorm2d = torch.nn.BatchNorm2d
model = eval('models.' + config.MODEL.NAME + '.get_seg_model')(config)
dump_input = torch.rand(
(1, 3, config.TRAIN.IMAGE_SIZE[1], config.TRAIN.IMAGE_SIZE[0]))
logger.info(get_model_summary(model.cuda(), dump_input.cuda()))
if config.TEST.MODEL_FILE:
model_state_file = config.TEST.MODEL_FILE
else:
model_state_file = os.path.join(
final_output_dir, 'best.pth') # turn final_state.pth to best.pth
logger.info('=> loading model from {}'.format(model_state_file))
pretrained_dict = torch.load(model_state_file)
if 'state_dict' in pretrained_dict:
pretrained_dict = pretrained_dict['state_dict']
model_dict = model.state_dict()
pretrained_dict = {
k[6:]: v
for k, v in pretrained_dict.items() if k[6:] in model_dict.keys()
}
for k, _ in pretrained_dict.items():
logger.info('=> loading {} from pretrained model'.format(k))
model_dict.update(pretrained_dict)
model.load_state_dict(model_dict)
gpus = list(config.GPUS)
model = nn.DataParallel(model, device_ids=gpus).cuda()
# prepare data
test_size = (config.TEST.IMAGE_SIZE[1], config.TEST.IMAGE_SIZE[0])
test_dataset = eval('datasets.' + config.DATASET.DATASET)(
root=config.DATASET.ROOT,
list_path=config.DATASET.TEST_SET,
num_samples=None,
num_classes=config.DATASET.NUM_CLASSES,
multi_scale=False,
flip=False,
ignore_label=config.TRAIN.IGNORE_LABEL,
base_size=config.TEST.BASE_SIZE,
crop_size=test_size,
downsample_rate=1)
testloader = torch.utils.data.DataLoader(test_dataset,
batch_size=1,
shuffle=False,
num_workers=config.WORKERS,
pin_memory=True)
start = timeit.default_timer()
if 'val' in config.DATASET.TEST_SET:
# mean_IoU, IoU_array, pixel_acc, mean_acc = testval(config,
# test_dataset,
# testloader,
# model)
#
# msg = 'MeanIU: {: 4.4f}, Pixel_Acc: {: 4.4f}, \
# Mean_Acc: {: 4.4f}, Class IoU: '.format(mean_IoU,
# pixel_acc, mean_acc)
# logging.info(msg)
# logging.info(IoU_array)
F, J = myTestval(config, test_dataset, testloader, model)
msg = 'F: {: 4.4f}, J: {: 4.4f} '.format(F, J)
logging.info(msg)
elif 'test' in config.DATASET.TEST_SET:
# import pdb
# pdb.set_trace()
test(config,
test_dataset,
testloader,
model,
sv_dir=r'/raid/wj/HRNet-Semantic-Segmentation/results/faceparse')
end = timeit.default_timer()
logger.info('Mins: %d' % np.int((end - start) / 60))
logger.info('Done')
def main():
args = parse_args()
logger, final_output_dir, tb_log_dir = create_logger(
config, args.cfg, 'train')
logger.info(pprint.pformat(args))
logger.info(config)
writer_dict = {
'writer': SummaryWriter(tb_log_dir),
'train_global_steps': 0,
'valid_global_steps': 0,
}
# cudnn related setting
cudnn.benchmark = config.CUDNN.BENCHMARK
cudnn.deterministic = config.CUDNN.DETERMINISTIC
cudnn.enabled = config.CUDNN.ENABLED
gpus = list(config.GPUS)
# build model
model = eval('models.' + config.MODEL.NAME + '.get_seg_model')(config)
dump_input = torch.rand(
(1, 3, config.TRAIN.IMAGE_SIZE[1], config.TRAIN.IMAGE_SIZE[0]))
logger.info(get_model_summary(model.cuda(), dump_input.cuda()))
# copy model file
this_dir = os.path.dirname(__file__)
models_dst_dir = os.path.join(final_output_dir, 'models')
if os.path.exists(models_dst_dir):
shutil.rmtree(models_dst_dir)
shutil.copytree(os.path.join(this_dir, '../lib/models'), models_dst_dir)
# prepare data
crop_size = (config.TRAIN.IMAGE_SIZE[1], config.TRAIN.IMAGE_SIZE[0])
train_dataset = eval('datasets.' + config.DATASET.DATASET)(
root=config.DATASET.ROOT,
list_path=config.DATASET.TRAIN_SET,
num_samples=None,
num_classes=config.DATASET.NUM_CLASSES,
multi_scale=config.TRAIN.MULTI_SCALE,
flip=config.TRAIN.FLIP,
ignore_label=config.TRAIN.IGNORE_LABEL,
base_size=config.TRAIN.BASE_SIZE,
crop_size=crop_size,
downsample_rate=config.TRAIN.DOWNSAMPLERATE,
scale_factor=config.TRAIN.SCALE_FACTOR)
trainloader = torch.utils.data.DataLoader(
train_dataset,
batch_size=config.TRAIN.BATCH_SIZE_PER_GPU * len(gpus),
shuffle=config.TRAIN.SHUFFLE,
num_workers=config.WORKERS,
pin_memory=True,
drop_last=True)
test_size = (config.TEST.IMAGE_SIZE[1], config.TEST.IMAGE_SIZE[0])
test_dataset = eval('datasets.' + config.DATASET.DATASET)(
root=config.DATASET.ROOT,
list_path=config.DATASET.TEST_SET,
num_samples=config.TEST.NUM_SAMPLES,
num_classes=config.DATASET.NUM_CLASSES,
multi_scale=False,
flip=False,
ignore_label=config.TRAIN.IGNORE_LABEL,
base_size=config.TEST.BASE_SIZE,
crop_size=test_size,
downsample_rate=1)
testloader = torch.utils.data.DataLoader(
test_dataset,
batch_size=config.TEST.BATCH_SIZE_PER_GPU * len(gpus),
shuffle=False,
num_workers=config.WORKERS,
pin_memory=True)
# criterion
if config.LOSS.USE_OHEM:
criterion = OhemCrossEntropy(ignore_label=config.TRAIN.IGNORE_LABEL,
thres=config.LOSS.OHEMTHRES,
min_kept=config.LOSS.OHEMKEEP,
weight=train_dataset.class_weights)
else:
criterion = CrossEntropy(ignore_label=config.TRAIN.IGNORE_LABEL,
weight=train_dataset.class_weights)
model = FullModel(model, criterion)
model = nn.DataParallel(model, device_ids=gpus).cuda()
# optimizer
if config.TRAIN.OPTIMIZER == 'sgd':
optimizer = torch.optim.SGD(
[{
'params': filter(lambda p: p.requires_grad,
model.parameters()),
'lr': config.TRAIN.LR
}],
lr=config.TRAIN.LR,
momentum=config.TRAIN.MOMENTUM,
weight_decay=config.TRAIN.WD,
nesterov=config.TRAIN.NESTEROV,
)
else:
raise ValueError('Only Support SGD optimizer')
epoch_iters = np.int(train_dataset.__len__() /
config.TRAIN.BATCH_SIZE_PER_GPU / len(gpus))
model_state_file = os.path.join(final_output_dir, 'best.pth')
if os.path.isfile(model_state_file):
checkpoint = torch.load(model_state_file)
best_mIoU = checkpoint['best_mIoU']
last_epoch = checkpoint['epoch']
model.module.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
logger.info("=> loaded checkpoint (epoch {})".format(
checkpoint['epoch']))
epoch = 0
valid_loss, mean_IoU = validate(config, testloader, model, writer_dict,
epoch)
msg = 'Validation Loss: {:.3f}, MeanIU: {: 4.4f}, Best_mIoU: {: 4.4f}'.format(
valid_loss, mean_IoU, best_mIoU)
logging.info(msg)
def main():
args = parse_args()
logger, final_output_dir, _ = create_logger(config, args.cfg,
'vis_gradcam')
logger.info(pprint.pformat(args))
logger.info(pprint.pformat(config))
# cudnn related setting
cudnn.benchmark = config.CUDNN.BENCHMARK
cudnn.deterministic = config.CUDNN.DETERMINISTIC
cudnn.enabled = config.CUDNN.ENABLED
# build model
model = eval('models.' + config.MODEL.NAME + '.get_seg_model')(config)
dump_input = torch.rand(
(1, 3, config.TRAIN.IMAGE_SIZE[1], config.TRAIN.IMAGE_SIZE[0]))
logger.info(get_model_summary(model.cuda(), dump_input.cuda()))
if config.TEST.MODEL_FILE:
model_state_file = config.TEST.MODEL_FILE
else:
model_state_file = os.path.join(final_output_dir, 'best.pth')
logger.info('=> loading model from {}'.format(model_state_file))
pretrained_dict = torch.load(model_state_file)
model_dict = model.state_dict()
pretrained_dict = {
k[6:]: v
for k, v in pretrained_dict.items() if k[6:] in model_dict.keys()
}
for k, _ in pretrained_dict.items():
logger.info('=> loading {} from pretrained model'.format(k))
model_dict.update(pretrained_dict)
model.load_state_dict(model_dict)
# Wrap original model with NBDT
if config.NBDT.USE_NBDT:
from nbdt.model import SoftSegNBDT
model = SoftSegNBDT(config.NBDT.DATASET,
model,
hierarchy=config.NBDT.HIERARCHY,
classes=class_names)
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = model.to(device).eval()
# Retrieve input image corresponding to args.image_index
test_size = (config.TEST.IMAGE_SIZE[1], config.TEST.IMAGE_SIZE[0])
test_dataset = eval('datasets.' + config.DATASET.DATASET)(
root=config.DATASET.ROOT,
list_path=config.DATASET.TEST_SET,
num_samples=None,
num_classes=config.DATASET.NUM_CLASSES,
multi_scale=False,
flip=False,
ignore_label=config.TRAIN.IGNORE_LABEL,
base_size=config.TEST.BASE_SIZE,
crop_size=test_size,
downsample_rate=1)
# Define target layer as final convolution layer if not specified
if args.target_layers:
target_layers = args.target_layers.split(',')
else:
for name, module in list(model.named_modules())[::-1]:
if isinstance(module, nn.Conv2d):
target_layers = [name]
break
logger.info('Target layers set to {}'.format(str(target_layers)))
# Append model. to target layers if using nbdt
if config.NBDT.USE_NBDT:
target_layers = ['model.' + layer for layer in target_layers]
def generate_and_save_saliency(image_index,
pixel_i=None,
pixel_j=None,
crop_size=None,
normalize=False):
"""too lazy to move out to global lol"""
nonlocal maximum, minimum, label
# Generate GradCAM + save heatmap
heatmaps = []
raw_image = retrieve_raw_image(test_dataset, image_index)
should_crop = crop_size is not None and pixel_i is not None and pixel_j is not None
if should_crop:
raw_image = crop(pixel_i,
pixel_j,
crop_size,
raw_image,
is_tensor=False)
for layer in target_layers:
gradcam_region = gradcam.generate(target_layer=layer,
normalize=False)
if should_crop:
gradcam_region = crop(pixel_i,
pixel_j,
crop_size,
gradcam_region,
is_tensor=True)
maximum = max(float(gradcam_region.max()), maximum)
minimum = min(float(gradcam_region.min()), minimum)
logger.info(f'=> Bounds: ({minimum}, {maximum})')
heatmaps.append(gradcam_region)
output_dir = generate_output_dir(final_output_dir, args.vis_mode,
layer, config.NBDT.USE_NBDT,
nbdt_node_wnid, args.crop_size,
args.nbdt_node_wnids_for)
save_path = generate_save_path(output_dir, gradcam_kwargs)
logger.info('Saving {} heatmap at {}...'.format(
args.vis_mode, save_path))
if normalize:
gradcam_region = GradCAM.normalize(gradcam_region)
save_gradcam(save_path,
gradcam_region,
raw_image,
save_npy=not args.skip_save_npy)
else:
save_gradcam(save_path,
gradcam_region,
raw_image,
minimum=minimum,
maximum=maximum,
save_npy=not args.skip_save_npy)
output_dir_original = output_dir + '_original'
os.makedirs(output_dir_original, exist_ok=True)
save_path_original = generate_save_path(output_dir_original,
gradcam_kwargs,
ext='jpg')
logger.info('Saving {} original at {}...'.format(
args.vis_mode, save_path_original))
cv2.imwrite(save_path_original, raw_image)
if crop_size and pixel_i and pixel_j:
continue
output_dir += '_overlap'
os.makedirs(output_dir, exist_ok=True)
save_path_overlap = generate_save_path(output_dir,
gradcam_kwargs,
ext='npy')
save_path_plot = generate_save_path(output_dir,
gradcam_kwargs,
ext='jpg')
if not args.skip_save_npy:
logger.info('Saving {} overlap data at {}...'.format(
args.vis_mode, save_path_overlap))
logger.info('Saving {} overlap plot at {}...'.format(
args.vis_mode, save_path_plot))
save_overlap(save_path_overlap,
save_path_plot,
gradcam_region,
label,
save_npy=not args.skip_save_npy)
if len(heatmaps) > 1:
combined = torch.prod(torch.stack(heatmaps, dim=0), dim=0)
combined /= combined.max()
save_path = generate_save_path(final_output_dir, args.vis_mode,
gradcam_kwargs, 'combined',
config.NBDT.USE_NBDT,
nbdt_node_wnid)
logger.info('Saving combined {} heatmap at {}...'.format(
args.vis_mode, save_path))
save_gradcam(save_path, combined, raw_image)
nbdt_node_wnids = args.nbdt_node_wnid or []
cls = args.nbdt_node_wnids_for
if cls:
assert config.NBDT.USE_NBDT, 'Must be using NBDT'
from nbdt.data.custom import Node
assert hasattr(model, 'rules') and hasattr(model.rules, 'nodes'), \
'NBDT must have rules with nodes'
logger.info("Getting nodes leading up to class leaf {}...".format(cls))
leaf_to_path_nodes = Node.get_leaf_to_path(model.rules.nodes)
cls_index = class_names.index(cls)
leaf = model.rules.nodes[0].wnids[cls_index]
path_nodes = leaf_to_path_nodes[leaf]
nbdt_node_wnids = [
item['node'].wnid for item in path_nodes if item['node']
]
def run():
nonlocal maximum, minimum, label, gradcam_kwargs
for image_index in get_image_indices(args.image_index,
args.image_index_range):
image, label, _, name = test_dataset[image_index]
image = torch.from_numpy(image).unsqueeze(0).to(device)
logger.info("Using image {}...".format(name))
pred_probs, pred_labels = gradcam.forward(image)
maximum, minimum = -1000, 0
logger.info(f'=> Starting bounds: ({minimum}, {maximum})')
if args.crop_for and class_names.index(args.crop_for) not in label:
print(
f'Skipping image {image_index} because no {args.crop_for} found'
)
continue
if getattr(Saliency, 'whole_image', False):
assert not (
args.pixel_i or args.pixel_j or args.pixel_i_range
or args.pixel_j_range), \
'the "Whole" saliency method generates one map for the whole ' \
'image, not for specific pixels'
gradcam_kwargs = {'image': image_index}
if args.suffix:
gradcam_kwargs['suffix'] = args.suffix
gradcam.backward(pred_labels[:, [0], :, :])
generate_and_save_saliency(image_index)
if args.crop_size <= 0:
continue
if args.crop_for:
cls_index = class_names.index(args.crop_for)
label = torch.Tensor(label).to(pred_labels.device)
# is_right_class = pred_labels[0,0,:,:] == cls_index
# is_correct = pred_labels == label
is_right_class = is_correct = label == cls_index #TODO:tmp
pixels = (is_right_class * is_correct).nonzero()
pixels = get_random_pixels(args.pixel_max_num_random,
pixels,
seed=cls_index)
else:
assert (args.pixel_i
or args.pixel_i_range) and (args.pixel_j
or args.pixel_j_range)
pixels = get_pixels(args.pixel_i, args.pixel_j,
args.pixel_i_range, args.pixel_j_range,
args.pixel_cartesian_product)
logger.info(f'Running on {len(pixels)} pixels.')
for pixel_i, pixel_j in pixels:
pixel_i, pixel_j = int(pixel_i), int(pixel_j)
assert pixel_i < test_size[0] and pixel_j < test_size[1], \
"Pixel ({},{}) is out of bounds for image of size ({},{})".format(
pixel_i,pixel_j,test_size[0],test_size[1])
# Run backward pass
# Note: Computes backprop wrt most likely predicted class rather than gt class
gradcam_kwargs = {
'image': image_index,
'pixel_i': pixel_i,
'pixel_j': pixel_j
}
if args.suffix:
gradcam_kwargs['suffix'] = args.suffix
logger.info(
f'Running {args.vis_mode} on image {image_index} at pixel ({pixel_i},{pixel_j}). Using filename suffix: {args.suffix}'
)
output_pixel_i, output_pixel_j = compute_output_coord(
pixel_i, pixel_j, test_size, pred_probs.shape[2:])
if not getattr(Saliency, 'whole_image', False):
gradcam.backward(pred_labels[:, [0], :, :], output_pixel_i,
output_pixel_j)
if args.crop_size <= 0:
generate_and_save_saliency(image_index)
else:
generate_and_save_saliency(image_index, pixel_i, pixel_j,
args.crop_size)
logger.info(f'=> Final bounds are: ({minimum}, {maximum})')
# Instantiate wrapper once, outside of loop
Saliency = METHODS[args.vis_mode]
gradcam = Saliency(model=model,
candidate_layers=target_layers,
use_nbdt=config.NBDT.USE_NBDT,
nbdt_node_wnid=None)
maximum, minimum, label, gradcam_kwargs = -1000, 0, None, {}
for nbdt_node_wnid in nbdt_node_wnids:
if config.NBDT.USE_NBDT:
logger.info("Using logits from node with wnid {}...".format(
nbdt_node_wnid))
gradcam.set_nbdt_node_wnid(nbdt_node_wnid)
run()
if not nbdt_node_wnids:
nbdt_node_wnid = None
run()
def 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 setting
cudnn.benchmark = config.CUDNN.BENCHMARK
torch.backends.cudnn.deterministic = config.CUDNN.DETERMINISTIC
torch.backends.cudnn.enabled = config.CUDNN.ENABLED
dataset = config.DATASET.DATASET
if dataset.startswith('imagenet'):
model = eval('models.' + config.MODEL.NAME + '.get_cls_net')(config)
else:
model = eval('models.' + config.MODEL.NAME +
'.get_cls_net_cifar')(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))
# n_flops, n_params = measure_model(model, 32, 32, torch.zeros(1,3,32,32))
# logger.info("param size = %fMB", n_params[0]/1e6)
# logger.info("flops = %fM", n_flops[0]/1e6)
# copy model file
this_dir = os.path.dirname(__file__)
models_dst_dir = os.path.join(final_output_dir, 'models')
if os.path.exists(models_dst_dir):
shutil.rmtree(models_dst_dir)
shutil.copytree(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 = torch.nn.DataParallel(model, device_ids=gpus).cuda()
# define loss function (criterion) and optimizer
criterion = torch.nn.CrossEntropyLoss().cuda()
optimizer = get_optimizer(config, model)
best_perf = 0.0
best_model = False
last_epoch = config.TRAIN.BEGIN_EPOCH
if 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.module.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
logger.info("=> loaded checkpoint (epoch {})".format(
checkpoint['epoch']))
best_model = True
if 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 code
if dataset.startswith('imagenet'):
traindir = 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,
]))
train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=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=config.TEST.BATCH_SIZE_PER_GPU * len(gpus),
shuffle=False,
num_workers=config.WORKERS,
pin_memory=True)
else:
CIFAR_MEAN = [0.49139968, 0.48215827, 0.44653124]
CIFAR_STD = [0.24703233, 0.24348505, 0.26158768]
if config.TRAIN.AUGMENT == 'autoaugment':
print("==>use autoaugment")
train_transform = transforms.Compose([
transforms.RandomCrop(32, padding=4, fill=128),
transforms.RandomHorizontalFlip(),
CIFAR10Policy(),
transforms.ToTensor(),
Cutout_v2(n_holes=1, length=16),
transforms.Normalize(CIFAR_MEAN, CIFAR_STD),
])
else:
train_transform = transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(CIFAR_MEAN, CIFAR_STD),
])
if config.TRAIN.AUGMENT == 'cutout':
train_transform.transforms.append(Cutout(16))
valid_transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(CIFAR_MEAN, CIFAR_STD),
])
train_dataset = datasets.CIFAR10(root=config.DATASET.ROOT,
train=True,
download=True,
transform=train_transform)
valid_dataset = datasets.CIFAR10(root=config.DATASET.ROOT,
train=False,
download=True,
transform=valid_transform)
train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=config.TRAIN.BATCH_SIZE_PER_GPU * len(gpus),
shuffle=True,
num_workers=config.WORKERS,
pin_memory=True)
valid_loader = torch.utils.data.DataLoader(
valid_dataset,
batch_size=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 epoch
train(config, train_loader, model, criterion, optimizer, epoch,
final_output_dir, tb_log_dir, writer_dict)
# evaluate on validation set
perf_indicator = validate(config, valid_loader, model, criterion,
final_output_dir, tb_log_dir, writer_dict)
if perf_indicator > best_perf:
best_perf = perf_indicator
best_model = True
else:
best_model = False
logger.info('=> saving checkpoint to {}'.format(final_output_dir))
save_checkpoint(
{
'epoch': epoch + 1,
'model': config.MODEL.NAME,
'state_dict': model.module.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.module.state_dict(), final_model_state_file)
writer_dict['writer'].close()
def main():
mean = [0.485, 0.456, 0.406],
std = [0.229, 0.224, 0.225]
args = parse_args()
logger, final_output_dir, _ = create_logger(config, args.cfg, 'test')
logger.info(pprint.pformat(args))
logger.info(pprint.pformat(config))
# cudnn related setting
cudnn.benchmark = config.CUDNN.BENCHMARK
cudnn.deterministic = config.CUDNN.DETERMINISTIC
cudnn.enabled = config.CUDNN.ENABLED
# build model
if torch.__version__.startswith('1'):
module = eval('models.' + config.MODEL.NAME)
module.BatchNorm2d_class = module.BatchNorm2d = torch.nn.BatchNorm2d
model = eval('models.' + config.MODEL.NAME + '.get_seg_model')(config)
dump_input = torch.rand(
(1, 3, config.TRAIN.IMAGE_SIZE[1], config.TRAIN.IMAGE_SIZE[0]))
logger.info(get_model_summary(model.cuda(), dump_input.cuda()))
if config.TEST.MODEL_FILE:
model_state_file = config.TEST.MODEL_FILE
else:
# model_state_file = os.path.join(final_output_dir, 'best_0.7589.pth')
model_state_file = os.path.join(final_output_dir, 'best.pth')
logger.info('=> loading model from {}'.format(model_state_file))
pretrained_dict = torch.load(
'/home/kong/Documents/DDRNet.Pytorch/DDRNet.Pytorch/output/face/ddrnet23_slim/checkpoint.pth.tar'
)
if 'state_dict' in pretrained_dict:
pretrained_dict = pretrained_dict['state_dict']
# print(pretrained_dict.keys())
new_state = {
k: v
for k, v in pretrained_dict.items() if k in model.state_dict()
}
model.load_state_dict(new_state)
model = model.cuda()
torch.save(model.state_dict(),
'model_best_bacc.pth.tar',
_use_new_zipfile_serialization=False)
# gpus = list(config.GPUS)
# model = nn.DataParallel(model, device_ids=gpus).cuda()
# prepare data
test_size = (config.TEST.IMAGE_SIZE[1], config.TEST.IMAGE_SIZE[0])
print(test_size)
img = cv2.imread(
"/home/kong/Documents/DDRNet.Pytorch/DDRNet.Pytorch/images/418cd0c0b416d93bc5a129834537f2e1.jpeg"
)
img = cv2.resize(img, (512, 512))
image = img.astype(np.float32)[:, :, ::-1]
image = image / 255.0
image -= mean
image /= std
image = image.transpose((2, 0, 1))
image = torch.from_numpy(image)
# image = image.permute((2, 0, 1))
# print(image.shape)
image = image.unsqueeze(0)
image = image.cuda()
start = time.time()
# print("input : ",image)
for i in range(1):
out = model(image)
end = time.time()
#print("Cuda 1000 images inference time : ",1000.0/(end - start))
outadd = out[0] * 1.0 + out[1] * 0.4
out0 = out[0].squeeze(dim=0)
out1 = out[1].squeeze(dim=0)
# print(out0.size(),out0[0,1,1],out0[1,1,1])
# print("out:",out0)
outadd = outadd.squeeze(dim=0)
out0 = F.softmax(out0, dim=0)
out1 = F.softmax(out1, dim=0)
outadd = F.softmax(outadd, dim=0)
out0 = torch.argmax(out0, dim=0)
out1 = torch.argmax(out1, dim=0)
outadd = torch.argmax(outadd, dim=0)
pred0 = out0.detach().cpu().numpy()
pred1 = out1.detach().cpu().numpy()
predadd = outadd.detach().cpu().numpy()
pred0 = pred0 * 255
pred1 = pred1 * 255
predadd = predadd * 255
####================= alpha channel =========================#
print("pred0:", pred0.shape, img.shape)
pred0 = np.array(pred0, np.uint8)
pred0up = cv2.resize(pred0, (512, 512))
png = np.dstack((img, pred0up))
bg = cv2.imread(
"/home/kong/Downloads/4aeb26a89778f73261ccef283e70992f.jpeg")
addpng = merge_img(bg, png, 500, 1012, 100, 612)
cv2.imwrite("png.png", addpng)
pred_ch = np.zeros(pred0.shape)
pred_rgb0 = np.array([pred_ch, pred_ch, pred0])
pred_rgb1 = np.array([pred_ch, pred_ch, pred1])
pred_rgbadd = np.array([predadd, pred_ch, predadd])
pred_rgb0 = pred_rgb0.transpose(1, 2, 0)
pred_rgb1 = pred_rgb1.transpose(1, 2, 0)
pred_rgbadd = pred_rgbadd.transpose(1, 2, 0)
pred_rgb0 = cv2.resize(pred_rgb0, (img.shape[1], img.shape[0]))
pred_rgb1 = cv2.resize(pred_rgb1, (img.shape[1], img.shape[0]))
pred_rgbadd = cv2.resize(pred_rgbadd, (img.shape[1], img.shape[0]))
dst = cv2.addWeighted(img, 0.7, pred_rgb0.astype(np.uint8), 0.3, 0)
dst1 = cv2.addWeighted(img, 0.7, pred_rgb1.astype(np.uint8), 0.3, 0)
dstadd = cv2.addWeighted(img, 0.7, pred_rgbadd.astype(np.uint8), 0.3, 0)
imgadd = np.hstack((img, dstadd))
cv2.imwrite("a242.jpg", imgadd)
def main():
args = parse_args()
logger, final_output_dir, tb_log_dir = create_logger(
config, args.cfg, 'train')
model_func = 'models.{0}.get_{0}'.format(config.MODEL.NAME)
model = eval(model_func)(config, pretrained=True)
dump_input = torch.rand((1, 3, config.MODEL.IMAGE_SIZE[0], config.MODEL.IMAGE_SIZE[1]))
print(get_model_summary(model, dump_input))
gpus = list(config.GPUS)
device = torch.device('cuda:{}'.format(gpus[0]) if torch.cuda.is_available() else 'cpu')
model.to(device)
optimizer = get_optimizer(config, model)
if config.FP16:
model, optimizer = amp.initialize(model, optimizer, opt_level='O1')
criterion = {'reid': nn.CrossEntropyLoss(), 'sr': nn.MSELoss()}
if config.TRAIN.SR_FILTER:
criterion['sr'] = MSEFilterLoss()
best_perf = 0.0
best_model = False
last_epoch = config.TRAIN.BEGIN_EPOCH
if 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'])
optimizer.load_state_dict(checkpoint['optimizer'])
if config.FP16:
amp.load_state_dict(checkpoint['amp'])
logger.info("=> loaded checkpoint (epoch {})"
.format(checkpoint['epoch']))
best_model = True
writer_dict = {
'writer': SummaryWriter(log_dir=tb_log_dir),
'train_global_steps': 0,
'valid_global_steps': 0,
}
train_dataset = ImagePersonDataset(config, mode='train')
train_sampler = RandomIdentitySampler(
train_dataset,
batch_size=config.TRAIN.BATCH_SIZE_PER_GPU * len(gpus),
num_instances=4
)
train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=config.TRAIN.BATCH_SIZE_PER_GPU * len(gpus),
sampler=train_sampler,
num_workers=config.WORKERS,
pin_memory=True
)
valid_loader = torch.utils.data.DataLoader(
ImagePersonDataset(config, mode='val'),
batch_size=config.TEST.BATCH_SIZE_PER_GPU * len(gpus),
num_workers=config.WORKERS,
pin_memory=True
)
since = time.time()
for epoch in range(last_epoch, config.TRAIN.END_EPOCH):
print('Epoch {}/{}'.format(epoch + 1, config.TRAIN.END_EPOCH))
train(config, train_loader, model, criterion, optimizer, epoch, device,
writer_dict)
perf_indicator = validate(config, valid_loader, model, criterion, device,
writer_dict)
if perf_indicator > best_perf:
best_perf = perf_indicator
best_model = True
else:
best_model = False
logger.info('=> saving checkpoint to {}'.format(final_output_dir))
checkpoint = {
'epoch': epoch + 1,
'model': config.MODEL.NAME,
'state_dict': model.state_dict(),
'perf': perf_indicator,
'optimizer': optimizer.state_dict()
}
if config.FP16:
checkpoint['amp'] = amp.state_dict()
save_checkpoint(checkpoint, best_model, final_output_dir,
filename='checkpoint.pth.tar')
h, m, s = get_hms_from_sec(time.time() - since)
print('=> total training time: {}h {}m {}s'.format(h, m, s))
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)
writer_dict['writer'].close()
def main():
mean = [0.485, 0.456, 0.406],
std = [0.229, 0.224, 0.225]
args = parse_args()
logger, final_output_dir, _ = create_logger(config, args.cfg, 'test')
logger.info(pprint.pformat(args))
logger.info(pprint.pformat(config))
# cudnn related setting
cudnn.benchmark = config.CUDNN.BENCHMARK
cudnn.deterministic = config.CUDNN.DETERMINISTIC
cudnn.enabled = config.CUDNN.ENABLED
# build model
if torch.__version__.startswith('1'):
module = eval('models.' + config.MODEL.NAME)
module.BatchNorm2d_class = module.BatchNorm2d = torch.nn.BatchNorm2d
model = eval('models.' + config.MODEL.NAME + '.get_seg_model')(config)
dump_input = torch.rand(
(1, 3, config.TRAIN.IMAGE_SIZE[1], config.TRAIN.IMAGE_SIZE[0]))
logger.info(get_model_summary(model.cuda(), dump_input.cuda()))
if config.TEST.MODEL_FILE:
model_state_file = config.TEST.MODEL_FILE
else:
# model_state_file = os.path.join(final_output_dir, 'best_0.7589.pth')
model_state_file = os.path.join(final_output_dir, 'best.pth')
logger.info('=> loading model from {}'.format(model_state_file))
pretrained_dict = torch.load(
'/home/hwits/Documents/CarVid/DDRNet/DDRNet.pytorch/model_best_bacc.pth.tar'
)
if 'state_dict' in pretrained_dict:
pretrained_dict = pretrained_dict['state_dict']
newstate_dict = {
k: v
for k, v in pretrained_dict.items() if k in model.state_dict()
}
# print(pretrained_dict.keys())
model.load_state_dict(newstate_dict)
model = model.cuda()
print(model)
x = torch.ones((1, 3, 1024, 1024)).cuda()
model_trt = torch2trt(model, [x])
y = model(x)
print(len(y))
y_trt = model_trt(x)
print(y_trt.shape)
print(len(y_trt))
# check the output against PyTorch
print(torch.max(torch.abs(y - y_trt)))
with open('ddrnet.engine', "wb") as f:
f.write(model_trt.engine.serialize())
def main():
args = parse_args()
logger, final_output_dir, _ = create_logger(config, args.cfg, 'vis_erf')
logger.info(pprint.pformat(args))
logger.info(pprint.pformat(config))
# cudnn related setting
cudnn.benchmark = config.CUDNN.BENCHMARK
cudnn.deterministic = config.CUDNN.DETERMINISTIC
cudnn.enabled = config.CUDNN.ENABLED
# build model
model = eval('models.' + config.MODEL.NAME + '.get_seg_model')(config)
dump_input = torch.rand(
(1, 3, config.TRAIN.IMAGE_SIZE[1], config.TRAIN.IMAGE_SIZE[0]))
logger.info(get_model_summary(model.cuda(), dump_input.cuda()))
if config.TEST.MODEL_FILE:
model_state_file = config.TEST.MODEL_FILE
else:
model_state_file = os.path.join(final_output_dir, 'best.pth')
logger.info('=> loading model from {}'.format(model_state_file))
pretrained_dict = torch.load(model_state_file)
model_dict = model.state_dict()
pretrained_dict = {
k[6:]: v
for k, v in pretrained_dict.items() if k[6:] in model_dict.keys()
}
for k, _ in pretrained_dict.items():
logger.info('=> loading {} from pretrained model'.format(k))
model_dict.update(pretrained_dict)
model.load_state_dict(model_dict)
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = model.to(device).eval()
# Generate input
logger.info('Generating input of ones...')
input = torch.ones((1, 3, config.TEST.IMAGE_SIZE[1],
config.TEST.IMAGE_SIZE[0])).float().to(device)
input = torch.autograd.Variable(input, requires_grad=True)
# Run input through network + compute backward pass
logger.info('Computing backward pass...')
output = model(input)
output_grad = torch.zeros_like(output)
center_i, center_j = output.shape[2] // 2, output.shape[3] // 2
erf_center_i, erf_center_j = center_i + args.offset_i, center_j + args.offset_j
output_grad[0, 0, erf_center_i, erf_center_j] = 1.
output.backward(gradient=output_grad)
input_grad = input.grad[0, 0].cpu().data.numpy()
# Compute receptive field rectangle
nonzero_i = np.nonzero(np.sum(input_grad, axis=1))
min_i, max_i = np.min(nonzero_i), np.max(nonzero_i)
nonzero_j = np.nonzero(np.sum(input_grad, axis=0))
min_j, max_j = np.min(nonzero_j), np.max(nonzero_j)
rect_origin = (min_j, min_i)
rect_h, rect_w = max_i - min_i, max_j - min_j
rect = patches.Rectangle(rect_origin,
rect_w,
rect_h,
linewidth=1,
edgecolor='g',
facecolor='none')
# Save plot of gradient
logger.info('Saving gradient map...')
fig = plt.figure()
ax = fig.add_subplot(111)
ax.imshow(input_grad, cmap='coolwarm')
ax.add_patch(rect)
title_values = [
output.shape[2], output.shape[3], erf_center_i, erf_center_j,
rect_origin[1], rect_origin[0], rect_h, rect_w
]
ax.set_title(
"Output Size: {} x {}, Output Pixel: ({},{})\n ERF Origin: ({},{}), ERF Size: {} x {}"
.format(*title_values))
save_path = os.path.join(
final_output_dir, 'erf_{}_{}.png'.format(args.offset_i, args.offset_j))
plt.savefig(save_path)
def main():
mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]
args = parse_args()
logger, final_output_dir, _ = create_logger(
config, args.cfg, 'test')
logger.info(pprint.pformat(args))
logger.info(pprint.pformat(config))
# cudnn related setting
cudnn.benchmark = config.CUDNN.BENCHMARK
cudnn.deterministic = config.CUDNN.DETERMINISTIC
cudnn.enabled = config.CUDNN.ENABLED
# build model
if torch.__version__.startswith('1'):
module = eval('models.'+config.MODEL.NAME)
module.BatchNorm2d_class = module.BatchNorm2d = torch.nn.BatchNorm2d
model = eval('models.'+config.MODEL.NAME +
'.get_seg_model')(config)
dump_input = torch.rand(
(1, 3, config.TRAIN.IMAGE_SIZE[1], config.TRAIN.IMAGE_SIZE[0])
)
logger.info(get_model_summary(model.cuda(), dump_input.cuda()))
if config.TEST.MODEL_FILE:
model_state_file = config.TEST.MODEL_FILE
else:
# model_state_file = os.path.join(final_output_dir, 'best_0.7589.pth')
model_state_file = os.path.join(final_output_dir, 'best.pth')
logger.info('=> loading model from {}'.format(model_state_file))
pretrained_dict = torch.load('/home/hwits/Documents/CarVid/DDRNet/segmentation/best_val_smaller.pth')
# print("pretrained_dict:",pretrained_dict.keys())
if 'state_dict' in pretrained_dict:
pretrained_dict = pretrained_dict['state_dict']
# print(pretrained_dict.keys())
new_state = {k[6:]:v for k,v in pretrained_dict.items() if k[6:] in model.state_dict()}
model.load_state_dict(new_state)
model = model.cuda()
torch.save(model.state_dict(), 'model_best_bacc.pth.tar', _use_new_zipfile_serialization=False)
# gpus = list(config.GPUS)
# model = nn.DataParallel(model, device_ids=gpus).cuda()
# prepare data
test_size = (config.TEST.IMAGE_SIZE[1], config.TEST.IMAGE_SIZE[0])
print(test_size)
img = cv2.imread("/home/hwits/berlin_000528_000019_leftImg8bit.png")
img = cv2.resize(img,(1024,1024))
image = img.astype(np.float32)[:, :, ::-1]
image = image / 255.0
image -= mean
image /= std
image = image.transpose((2,0,1))
image = torch.from_numpy(image)
# image = image.permute((2, 0, 1))
print(image.shape)
image = image.unsqueeze(0)
image = image.cuda()
start = time.time()
# print("input : ",image)
for i in range(1):
out= model(image)
end = time.time()
print("out : ",out[0].shape)
pred = F.interpolate(
out[0], (img.shape[0],img.shape[1]),
mode='nearest'
)
print("results : ",pred.shape)
_, pred = torch.max(pred, dim=1)
pred = pred.squeeze(0).cpu().numpy()
map16.visualize_result(img, pred, '.', 'cityscape.jpg')
print("Cuda 1000 images inference time : ",1000.0/(end - start))
outadd = out[0]*1.0 + out[1]*0.4
out0 = out[0].squeeze(dim=0)
out1 = out[1].squeeze(dim=0)
outadd = outadd.squeeze(dim=0)
out0 = F.softmax(out0,dim=0)
out1 = F.softmax(out1,dim=0)
outadd = F.softmax(outadd,dim=0)
out0 = torch.argmax(out0,dim=0)
out1 = torch.argmax(out1,dim=0)
outadd = torch.argmax(outadd,dim=0)
pred0 = out0.detach().cpu().numpy()
pred1 = out1.detach().cpu().numpy()
predadd = outadd.detach().cpu().numpy()
pred0 = pred0*255
pred1 = pred1*255
predadd = predadd*255
pred_ch = np.zeros(pred0.shape)
pred_rgb0 = np.array([pred_ch,pred_ch,pred0])
pred_rgb1 = np.array([pred_ch,pred_ch,pred1])
pred_rgbadd = np.array([predadd,pred_ch,predadd])
pred_rgb0 = pred_rgb0.transpose(1,2,0)
pred_rgb1 = pred_rgb1.transpose(1,2,0)
pred_rgbadd = pred_rgbadd.transpose(1,2,0)
pred_rgb0 = cv2.resize(pred_rgb0,(img.shape[1],img.shape[0]))
pred_rgb1 = cv2.resize(pred_rgb1,(img.shape[1],img.shape[0]))
pred_rgbadd = cv2.resize(pred_rgbadd,(img.shape[1],img.shape[0]))
dst=cv2.addWeighted(img,0.7,pred_rgb0.astype(np.uint8),0.3,0)
dst1=cv2.addWeighted(img,0.7,pred_rgb1.astype(np.uint8),0.3,0)
dstadd=cv2.addWeighted(img,0.7,pred_rgbadd.astype(np.uint8),0.3,0)
imgadd = np.vstack((img,pred_rgb0,dst,pred_rgb1, dst1,pred_rgbadd, dstadd))
cv2.imwrite("a22.jpg",imgadd)
def 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 setting
cudnn.benchmark = config.CUDNN.BENCHMARK
torch.backends.cudnn.deterministic = config.CUDNN.DETERMINISTIC
torch.backends.cudnn.enabled = config.CUDNN.ENABLED
if config.MODEL.DOWNSAMPLE:
model = eval('models.' + config.MODEL.NAME +
'.get_cls_net')(config).cuda()
else:
model = eval('models.' + config.MODEL.NAME +
'.get_linear_net')(config).cuda()
dump_input = torch.rand(
(1, 3, config.MODEL.IMAGE_SIZE[1], config.MODEL.IMAGE_SIZE[0])).cuda()
logger.info(get_model_summary(model, dump_input))
if config.TEST.MODEL_FILE:
logger.info('=> loading model from {}'.format(config.TEST.MODEL_FILE))
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 = torch.nn.DataParallel(model, device_ids=gpus).cuda()
# define loss function (criterion) and optimizer
criterion = torch.nn.CrossEntropyLoss().cuda()
# Data loading code
valdir = os.path.join(config.DATASET.ROOT + '/images',
config.DATASET.TEST_SET)
dataset_name = config.DATASET.DATASET
if dataset_name == 'imagenet':
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)
else:
normalize = transforms.Normalize((0.4914, 0.4822, 0.4465),
(0.2023, 0.1994, 0.2010))
transform_valid = transforms.Compose([
transforms.ToTensor(),
normalize,
])
if dataset_name == 'cifar10':
valid_dataset = datasets.CIFAR10(root=f'{config.DATASET.ROOT}',
train=False,
download=True,
transform=transform_valid)
elif dataset_name == 'cifar100':
valid_dataset = datasets.CIFAR100(root=f'{config.DATASET.ROOT}',
train=False,
download=True,
transform=transform_valid)
valid_loader = torch.utils.data.DataLoader(
valid_dataset,
batch_size=config.TEST.BATCH_SIZE_PER_GPU * len(gpus),
shuffle=False,
num_workers=config.WORKERS,
pin_memory=True)
# evaluate on validation set
validate(config,
valid_loader,
model,
criterion,
None,
-1,
final_output_dir,
tb_log_dir,
None,
topk=(1, 5))
def 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 setting
cudnn.benchmark = config.CUDNN.BENCHMARK
torch.backends.cudnn.deterministic = config.CUDNN.DETERMINISTIC
torch.backends.cudnn.enabled = config.CUDNN.ENABLED
model = 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 file
this_dir = os.path.dirname(__file__)
models_dst_dir = os.path.join(final_output_dir, 'models')
if os.path.exists(models_dst_dir):
shutil.rmtree(models_dst_dir)
shutil.copytree(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 = [0,1]
model = torch.nn.DataParallel(model, device_ids=gpus).cuda()
# define loss function (criterion) and optimizer
criterion = torch.nn.CrossEntropyLoss().cuda()
optimizer = get_optimizer(config, model)
best_perf = 0.0
best_model = False
last_epoch = config.TRAIN.BEGIN_EPOCH
if 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.module.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
logger.info("=> loaded checkpoint (epoch {})"
.format(checkpoint['epoch']))
best_model = True
if 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 code
traindir = os.path.join(config.DATASET.ROOT, config.DATASET.TRAIN_SET)
valdir = os.path.join(config.DATASET.ROOT, config.DATASET.TEST_SET)
train_dataset = datasets(dataset_root='./Data/',split='train',size= config.MODEL.IMAGE_SIZE[0])
test_dataset = datasets(dataset_root='./Data/',split='test',size= config.MODEL.IMAGE_SIZE[0])
train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=config.TRAIN.BATCH_SIZE_PER_GPU*len(gpus),
shuffle=True,
num_workers=config.WORKERS,
pin_memory=True
)
valid_loader = torch.utils.data.DataLoader(
test_dataset,
batch_size=config.TRAIN.BATCH_SIZE_PER_GPU*len(gpus),
shuffle=True,
num_workers=config.WORKERS,
pin_memory=True
)
for epoch in range(last_epoch, config.TRAIN.END_EPOCH):
lr_scheduler.step()
# train for one epoch
train(config, train_loader, model, criterion, optimizer, epoch,
final_output_dir, tb_log_dir, writer_dict)
# evaluate on validation set
perf_indicator = validate(config, valid_loader, model, criterion,
final_output_dir, tb_log_dir, writer_dict)
if perf_indicator > best_perf:
best_perf = perf_indicator
best_model = True
torch.save(best_model, './best-model.pth')
else:
best_model = False
logger.info('=> saving checkpoint to {}'.format(final_output_dir))
save_checkpoint({
'epoch': epoch + 1,
'model': config.MODEL.NAME,
'state_dict': model.module.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.module.state_dict(), final_model_state_file)
writer_dict['writer'].close()
def main():
args = parse_args()
logger, final_output_dir, _ = create_logger(config, args.cfg, 'test')
logger.info(pprint.pformat(args))
logger.info(pprint.pformat(config))
# cudnn related setting
cudnn.benchmark = config.CUDNN.BENCHMARK
cudnn.deterministic = config.CUDNN.DETERMINISTIC
cudnn.enabled = config.CUDNN.ENABLED
# build model
if torch.__version__.startswith('1'):
module = eval('models.' + config.MODEL.NAME)
module.BatchNorm2d_class = module.BatchNorm2d = torch.nn.BatchNorm2d
model = eval('models.' + config.MODEL.NAME + '.get_seg_model')(config)
dump_input = torch.rand(
(1, 3, config.TRAIN.IMAGE_SIZE[1], config.TRAIN.IMAGE_SIZE[0]))
logger.info(get_model_summary(model.cuda(), dump_input.cuda()))
if config.TEST.MODEL_FILE:
model_state_file = config.TEST.MODEL_FILE
else:
# model_state_file = os.path.join(final_output_dir, 'best_0.7589.pth')
model_state_file = os.path.join(final_output_dir, 'best.pth')
logger.info('=> loading model from {}'.format(model_state_file))
pretrained_dict = torch.load(model_state_file)
if 'state_dict' in pretrained_dict:
pretrained_dict = pretrained_dict['state_dict']
model_dict = model.state_dict()
pretrained_dict = {
k[6:]: v
for k, v in pretrained_dict.items() if k[6:] in model_dict.keys()
}
for k, _ in pretrained_dict.items():
logger.info('=> loading {} from pretrained model'.format(k))
model_dict.update(pretrained_dict)
model.load_state_dict(model_dict)
model = model.cuda()
# gpus = list(config.GPUS)
# model = nn.DataParallel(model, device_ids=gpus).cuda()
# prepare data
test_size = (config.TEST.IMAGE_SIZE[1], config.TEST.IMAGE_SIZE[0])
test_dataset = eval('datasets.' + config.DATASET.DATASET)(
root=config.DATASET.ROOT,
list_path=config.DATASET.TEST_SET,
num_samples=None,
num_classes=config.DATASET.NUM_CLASSES,
multi_scale=False,
flip=False,
ignore_label=config.TRAIN.IGNORE_LABEL,
base_size=config.TEST.BASE_SIZE,
crop_size=test_size,
downsample_rate=1)
testloader = torch.utils.data.DataLoader(test_dataset,
batch_size=1,
shuffle=False,
num_workers=config.WORKERS,
pin_memory=True)
start = timeit.default_timer()
test(config,
test_dataset,
testloader,
model,
sv_dir=final_output_dir + '/test_result')
end = timeit.default_timer()
logger.info('Mins: %d' % np.int((end - start) / 60))
logger.info('Done')
def main():
args = parse_args()
logger, final_output_dir, _ = create_logger(
config, args.cfg, 'test')
logger.info(pprint.pformat(args))
logger.info(pprint.pformat(config))
# cudnn related setting
cudnn.benchmark = config.CUDNN.BENCHMARK
cudnn.deterministic = config.CUDNN.DETERMINISTIC
cudnn.enabled = config.CUDNN.ENABLED
# build model
model = models.gsto_hrnet.get_seg_model(config)
dump_input = torch.rand(
(1, 3, config.TRAIN.IMAGE_SIZE[1], config.TRAIN.IMAGE_SIZE[0])
)
logger.info(get_model_summary(model.cuda(), dump_input.cuda()))
if config.TEST.MODEL_FILE:
model_state_file = config.TEST.MODEL_FILE
else:
model_state_file = os.path.join(final_output_dir,
'final_state.pth')
logger.info('=> loading model from {}'.format(model_state_file))
model_state_file = config.TEST.MODEL_FILE
pretrained_dict = torch.load(model_state_file)
model_dict = model.state_dict()
pretrained_dict = {k[6:]: v for k, v in pretrained_dict.items()
if k[6:] in model_dict.keys()}
model_dict.update(pretrained_dict)
model.load_state_dict(model_dict)
gpus = list(config.GPUS)
model = nn.DataParallel(model, device_ids=gpus).cuda()
# prepare data
test_size = (config.TEST.IMAGE_SIZE[1], config.TEST.IMAGE_SIZE[0])
test_dataset = eval('datasets.' + config.DATASET.DATASET)(
root=config.DATASET.ROOT,
list_path=config.DATASET.TEST_SET,
num_samples=None,
num_classes=config.DATASET.NUM_CLASSES,
multi_scale=False,
flip=False,
ignore_label=config.TRAIN.IGNORE_LABEL,
base_size=config.TEST.BASE_SIZE,
crop_size=test_size,
downsample_rate=1)
testloader = torch.utils.data.DataLoader(
test_dataset,
batch_size=1,
shuffle=False,
num_workers=config.WORKERS,
pin_memory=True)
start = timeit.default_timer()
if 'val' in config.DATASET.TEST_SET:
mean_IoU, IoU_array, pixel_acc, mean_acc = testval(config,
test_dataset,
testloader,
model)
msg = 'MeanIU: {: 4.4f}, Pixel_Acc: {: 4.4f}, \
Mean_Acc: {: 4.4f}, Class IoU: '.format(mean_IoU,
pixel_acc, mean_acc)
logging.info(msg)
logging.info(IoU_array)
elif 'test' in config.DATASET.TEST_SET:
test(config,
test_dataset,
testloader,
model,
sv_dir=final_output_dir)
end = timeit.default_timer()
logger.info('Mins: %d' % np.int((end - start) / 60))
logger.info('Done')
def __init__(self, config):
super().__init__(config)
logger.info(config)
# cudnn related setting
cudnn.benchmark = config.CUDNN.BENCHMARK
cudnn.deterministic = config.CUDNN.DETERMINISTIC
cudnn.enabled = config.CUDNN.ENABLED
gpus = list(config.GPUS)
self.writer_dict = {
'writer': SummaryWriter(config.SUMMARY_DIR),
'train_global_steps': 0,
'valid_global_steps': 0,
}
self.checkpoint_dir = config.CHECKPOINT_DIR
self.output_dir = config.OUTPUT_DIR
self.num_classes = config.DATASET.NUM_CLASSES
if config.MODEL.NAME == 'unet':
self.model = UNet(3, config.DATASET.NUM_CLASSES)
elif config.MODEL.NAME == 'dl_resnet':
self.model = DeepLabv3_plus_res(
3,
n_classes=config.DATASET.NUM_CLASSES,
os=config.MODEL.OS,
pretrained=True)
elif config.MODEL.NAME == 'dl_xception':
self.model = DeepLabv3_plus_xce(
3,
n_classes=config.DATASET.NUM_CLASSES,
os=config.MODEL.OS,
pretrained=True)
dump_input = torch.rand(
(1, 3, config.TRAIN.IMAGE_SIZE[1], config.TRAIN.IMAGE_SIZE[0]))
logger.info(get_model_summary(self.model.cuda(), dump_input.cuda()))
self.train_dataset = RssraiDataset('train',
config.DATASET,
mean=[0.2797, 0.3397, 0.3250],
std=[0.1271, 0.1446, 0.1320])
self.trainloader = torch.utils.data.DataLoader(
self.train_dataset,
batch_size=config.TRAIN.BATCH_SIZE_PER_GPU * len(gpus),
shuffle=config.TRAIN.SHUFFLE,
num_workers=config.WORKERS,
pin_memory=True,
drop_last=True)
self.test_dataset = RssraiDataset('val',
config.DATASET,
mean=[0.2797, 0.3397, 0.3250],
std=[0.1271, 0.1446, 0.1320])
self.testloader = torch.utils.data.DataLoader(
self.test_dataset,
batch_size=config.TEST.BATCH_SIZE_PER_GPU * len(gpus),
shuffle=False,
num_workers=config.WORKERS,
pin_memory=True)
if config.LOSS.USE_OHEM:
self.criterion = OhemCrossEntropy(thres=config.LOSS.OHEMTHRES,
min_kept=config.LOSS.OHEMKEEP)
else:
self.criterion = CrossEntropy()
if len(gpus) > 1:
self.model = nn.DataParallel(self.model, device_ids=gpus).cuda()
else:
self.model.cuda()
# optimizer
if config.TRAIN.OPTIMIZER == 'sgd':
self.optimizer = torch.optim.SGD(
[{
'params':
filter(lambda p: p.requires_grad, self.model.parameters()),
'lr':
config.TRAIN.LR
}],
lr=config.TRAIN.LR,
momentum=config.TRAIN.MOMENTUM,
weight_decay=config.TRAIN.WD,
nesterov=config.TRAIN.NESTEROV,
)
elif config.TRAIN.OPTIMIZER == 'adam':
self.optimizer = torch.optim.Adam([{
'params':
filter(lambda p: p.requires_grad, self.model.parameters()),
'lr':
config.TRAIN.LR
}],
lr=config.TRAIN.LR,
weight_decay=config.TRAIN.WD)
else:
raise ValueError('Only Support SGD optimizer')
self.lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
self.optimizer, config.TRAIN.LR_STEP, config.TRAIN.LR_FACTOR)
self.start_epoch = 0
self.current_epoch = self.start_epoch
self.end_epoch = config.TRAIN.EPOCH
init_model_weights(self.model)
if config.TRAIN.AUTO_RESUME:
self.load_checkpoint(
os.path.join(config.CHECKPOINT_DIR, 'checkpoint.pth.tar'))
self.best_mIoU = -1
