def main(aux=False):
mode = args.mode
model_name = args.model
dataloaders = load_data(args)
init_model = create_model(model_name,
pretrained=args.pre_trained,
num_classes=args.num_classes)
init_model = init_model.to(args.gpu)
optimizer_ft = optim.AdamW(filter(lambda p: p.requires_grad,
init_model.parameters()),
lr=0.0001)
criterion = torch.nn.CrossEntropyLoss()
if args.mode == "all":
data_list = ["train", "val_right", "val_left"]
else:
data_list = ["train", "val"]
print(model_name)
train_model(init_model,
dataloaders,
criterion,
optimizer_ft,
"../records/saved_model/" + model_name + "_" + mode + ".pt",
args=args,
data_list=data_list,
aux=aux)
def load_backbone(args):
bone = create_model(
args.model,
pretrained=args.pre_trained,
num_classes=args.num_classes)
if args.dataset == "MNIST":
bone.conv1 = nn.Conv2d(1, 64, 3, stride=2, padding=1, bias=False)
if args.use_slot:
if args.use_pre:
checkpoint = torch.load(f"saved_model/{args.dataset}_no_slot_checkpoint.pth")
new_state_dict = OrderedDict()
for k, v in checkpoint["model"].items():
name = k[9:] # remove `backbone.`
new_state_dict[name] = v
bone.load_state_dict(new_state_dict)
print("load pre dataset parameter over")
if not args.grad:
if 'seresnet' in args.model:
bone.avg_pool = Identical()
bone.last_linear = Identical()
elif 'res' in args.model:
bone.global_pool = Identical()
bone.fc = Identical()
elif 'efficient' in args.model:
bone.global_pool = Identical()
bone.classifier = Identical()
elif 'densenet' in args.model:
bone.global_pool = Identical()
bone.classifier = Identical()
elif 'mobilenet' in args.model:
bone.global_pool = Identical()
bone.conv_head = Identical()
bone.act2 = Identical()
bone.classifier = Identical()
return bone
def __init__(self, args, dictionary):
super().__init__(dictionary)
if 'custom_size' in args.deit_arch:
self.deit = create_model(args.deit_arch,
pretrained=True,
img_size=args.input_size,
ape=args.ape,
mask_ratio=args.mask_ratio)
else:
self.deit = create_model(args.deit_arch,
pretrained=True,
ape=args.ape,
mask_ratio=args.mask_ratio)
self.fp16 = args.fp16
def for_vis(args, cal_part):
model_name = args.model
model = create_model(
model_name,
pretrained=args.pre_trained,
num_classes=args.num_classes)
# for name, module in model._modules.items():
# print(name)
model.load_state_dict(torch.load("saved_model/" + model_name + "_" + args.mode + ".pt", map_location=args.gpu), strict=True)
model.to(args.gpu)
model.eval()
args.mode = cal_part
if test_one:
root = "/home/wbw/PAN/final_cxr/nodule/CR-01084483_20150513"
location = np.array([0,0,1,1,0,0,0,1])
name = root.split("/")[-2]
img = Image.open(root + "/" + args.mode + ".png")
mask = Image.open(root + "/" + args.mode + "_mask.png")
size = mask.size
att = cal(name, model, img, size, label_true=1)
cal_siou(name, img, att, root, location)
else:
if test:
files = MakeListTest(args).get_files()
circle(files, model)
else:
PP = MakeList(args, need_part=False)
PP.file_name = ["nodule"]
train, val = PP.get_wanted_data()
circle(train, model)
circle(val, model)
def __init__(self,
model_name,
detail=False,
device='cuda',
torchscript=False,
precision='float32',
fuser='',
num_warm_iter=10,
num_bench_iter=50,
use_train_size=False,
**kwargs):
self.model_name = model_name
self.detail = detail
self.device = device
self.use_amp, self.model_dtype, self.data_dtype = resolve_precision(
precision)
self.channels_last = kwargs.pop('channels_last', False)
self.amp_autocast = torch.cuda.amp.autocast if self.use_amp else suppress
if fuser:
set_jit_fuser(fuser)
self.model = create_model(
model_name,
num_classes=kwargs.pop('num_classes', None),
in_chans=3,
global_pool=kwargs.pop('gp', 'fast'),
scriptable=torchscript,
drop_rate=kwargs.pop('drop', 0.),
drop_path_rate=kwargs.pop('drop_path', None),
drop_block_rate=kwargs.pop('drop_block', None),
)
self.model.to(
device=self.device,
dtype=self.model_dtype,
memory_format=torch.channels_last if self.channels_last else None)
self.num_classes = self.model.num_classes
self.param_count = count_params(self.model)
_logger.info('Model %s created, param count: %d' %
(model_name, self.param_count))
self.scripted = False
if torchscript:
self.model = torch.jit.script(self.model)
self.scripted = True
data_config = resolve_data_config(kwargs,
model=self.model,
use_test_size=not use_train_size)
self.input_size = data_config['input_size']
self.batch_size = kwargs.pop('batch_size', 256)
self.example_inputs = None
self.num_warm_iter = num_warm_iter
self.num_bench_iter = num_bench_iter
self.log_freq = num_bench_iter // 5
if 'cuda' in self.device:
self.time_fn = partial(cuda_timestamp, device=self.device)
else:
self.time_fn = timestamp
def main():
global args
setup_default_logging()
args, args_text = _parse_args()
global alpha
alpha = args.alpha
assert alpha == 0.67749, 'Make sure alpha = 0.67749'
args.device = 'cuda:0'
model = create_model(args.model,
pretrained=args.pretrained,
num_classes=args.num_classes,
drop_rate=args.drop,
global_pool=args.gp,
bn_tf=False,
bn_momentum=None,
bn_eps=None,
checkpoint_path=args.initial_checkpoint)
assert args.binarizable == 'T', 'Make sure --binarizable T'
if args.binarizable:
Model_binary_patch(model)
model.cuda()
use_amp = False
if has_apex and args.amp:
print('Using amp with --opt-level {}.'.format(args.opt_level))
model = amp.initialize(model, opt_level=args.opt_level)
use_amp = True
else:
print('Do NOT use amp.')
data_config = resolve_data_config(vars(args),
model=model,
verbose=args.local_rank == 0)
dataset_eval = torchvision.datasets.CIFAR100(root='~/Downloads/CIFAR100',
train=False,
download=True)
loader_eval = create_loader_CIFAR100(
dataset_eval,
input_size=data_config['input_size'],
batch_size=4 * args.batch_size,
is_training=False,
use_prefetcher=False,
interpolation=None,
mean=data_config['mean'],
std=data_config['std'],
num_workers=args.workers,
distributed=False,
)
validate_loss_fn = nn.CrossEntropyLoss().cuda()
eval_metrics = validate(model, loader_eval, validate_loss_fn, args)
def create_vitstr(num_tokens, model=None, checkpoint_path=''):
vitstr = create_model(model,
pretrained=True,
num_classes=num_tokens,
checkpoint_path=checkpoint_path)
# might need to run to get zero init head for transfer learning
vitstr.reset_classifier(num_classes=num_tokens)
return vitstr
def get_model(args):
print(f"Creating model: {args.model}")
model = create_model(
args.model,
pretrained=False,
drop_path_rate=args.drop_path,
drop_block_rate=None,
use_shared_rel_pos_bias=args.rel_pos_bias,
use_abs_pos_emb=args.abs_pos_emb,
init_values=args.layer_scale_init_value,
)
return model
def get_model(args):
if args.model == 'resnet50':
model = models.resnet.resnet50(pretrained=args.pretrained)
elif args.model == 'resnet18':
model = models.resnet.resnet18(pretrained=args.pretrained)
elif args.model == 'resnet34':
model = models.resnet.resnet34(pretrained=args.pretrained)
elif args.model == 'resnet101':
model = models.resnet.resnet101(pretrained=args.pretrained)
elif args.model == 'resnet152':
model = models.resnet.resnet152(pretrained=args.pretrained)
elif args.model == 'wide_resnet50_2':
model = models.resnet.wide_resnet50_2(pretrained=args.pretrained)
elif args.model == 'wide_resnet101_2':
model = models.resnet.wide_resnet101_2(pretrained=args.pretrained)
elif args.model == 'googlenet':
model = models.googlenet.googlenet(pretrained=args.pretrained)
elif args.model == 'vgg16':
model = models.vgg.vgg16(pretrained=args.pretrained)
elif args.model == 'mobilenet_v2':
model = models.mobilenetv2.mobilenet_v2(pretrained=args.pretrained)
elif args.model == 'inceptionv3':
model = models.inceptionv3.inception_v3(pretrained=args.pretrained)
elif args.model == 'inceptionv4':
model = models.inceptionv4.inceptionv4(pretrained="imagenet")
elif args.model == 'inceptionresnetv2':
model = models.inception_resnetv2.inceptionresnetv2(
pretrained="imagenet")
elif args.model == 'densenet121':
model = models.densenet.densenet121(pretrained=args.pretrained)
elif args.model == 'densenet161':
model = models.densenet.densenet161(pretrained=args.pretrained)
elif args.model == 'densenet169':
model = models.densenet.densenet169(pretrained=args.pretrained)
elif args.model == 'densenet201':
model = models.densenet.densenet201(pretrained=args.pretrained)
elif args.model.split('_')[0] == 'vit':
model = create_model(args.model,
pretrained=True,
num_classes=1000,
in_chans=3,
global_pool=args.gp,
scriptable=args.torchscript)
return model
def __init__(
self,
in_channels: int,
num_classes: int,
pretrained=False,
**kwargs,
) -> None:
super().__init__()
self.save_hyperparameters()
self.best_train_metric = None
self.best_valid_metric = None
self.model = create_model(
model_name=self.hparams.arch,
pretrained=pretrained,
num_classes=num_classes,
in_chans=in_channels,
)
def main():
args = parser.parse_args()
setup_default_logging(log_path=args.log_path)
_logger = logging.getLogger('inference')
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# create model
model = create_model(args.model,
num_classes=args.num_classes,
in_chans=3,
checkpoint_path=args.checkpoint)
data_config = resolve_data_config(vars(args),
model=model) # 默认配置 mean, std
img_transform = create_transform(data_config['input_size'],
is_training=False,
mean=data_config['mean'],
std=data_config['std'],
crop_pct=data_config['crop_pct'])
model = model.to(device)
model.eval()
# confusion matrix
conf_matrix = np.zeros((args.num_classes, args.num_classes))
# predict
class_dir = os.listdir(args.data)
for cl in class_dir:
imgs = os.listdir(os.path.join(args.data, cl))
for img in imgs:
im = Image.open(os.path.join(args.data, cl, img))
im1 = img_transform(im).unsqueeze(0)
im1 = im1.to(device)
outputs = model(im1)
idx = torch.max(outputs, 1)[1].tolist()[0]
probs = torch.nn.functional.softmax(outputs, dim=1)[0].tolist()
_logger.info('%s: %d, %.2f' %
(os.path.join(args.data, cl, img), idx, probs[idx]))
conf_matrix[int(cl), idx] += 1
_logger.info(conf_matrix)
def model_init(config_files=None, checkpoint_files=None, classify_cfg=None):
if config_files == None:
config_files = CONFIG_FILE_GENERIC
if checkpoint_files == None:
checkpoint_files = CHECKPOINT_FILE_GENERIC
if classify_cfg == None:
classify_cfg = CLASSIFY_CFG
model_det = init_detector(config_files, checkpoint_files, device='cuda:0')
model_classify = create_model(classify_cfg['model'],
num_classes=classify_cfg['num_classes'],
in_chans=3,
checkpoint_path=classify_cfg['checkpoint'])
logging.info('Model %s created, param count: %d' %
(classify_cfg['model'], sum([m.numel() for m in model_classify.parameters()])))
model_classify = model_classify.cuda()
model_classify.eval()
return model_det, model_classify
def main():
model = create_model(
'resnet50',
pretrained=False,
num_classes=1000,
drop_rate=0.,
drop_path_rate=0.,
global_pool='avg',
bn_tf=False,
bn_momentum=None,
bn_eps=None,
checkpoint_path='',
use_se=True)
model_profiling(model, 224, 224, use_cuda=False)
flops = add_flops(model)
print('model.flops', flops)
def __init__(self,
name,
return_interm_layers,
main_layer=-1,
group_norm=False):
super().__init__()
backbone = create_model(name,
pretrained=True,
in_chans=3,
features_only=True,
out_indices=(1, 2, 3, 4))
with torch.no_grad():
replace_bn(backbone)
num_channels = backbone.feature_info.channels()[-1]
self.body = backbone
self.num_channels = num_channels
self.interm = return_interm_layers
self.main_layer = main_layer
def create_classification_model(self, model_name):
if str(model_name.value) == "efficientnet_b3a":
model_path = os.getenv("HOME") + '/ros2_models/effecient.pth'
model = create_model('efficientnet_b3a',
False,
model_path,
num_classes=6)
config = resolve_data_config({}, model=model)
self.transform = create_transform(**config)
# sd = torch.load(model_path, map_location=self.device)
# model = models.resnet18()
# model.fc = torch.nn.Linear(512, 6)
# model.load_state_dict(sd['state_dict'])
return model
print("Only valid model option currently is efficientnet_b3a")
def objective(trial):
torch.cuda.empty_cache()
device = torch.device(opt.device)
criterion = LabelSmoothingLoss(classes=43, smoothing=0.1)
drop_path_rate = trial.suggest_discrete_uniform(
'drop_path_rate', 0.0, 0.6, 0.1)
model = create_model(
model_name,
pretrained=True,
num_classes=43,
drop_rate=drop_path_rate,
drop_connect_rate=None, # DEPRECATED, use drop_path
drop_path_rate=None,
drop_block_rate=None,
global_pool=None,
bn_tf=False,
bn_momentum=None,
bn_eps=None,
checkpoint_path=None)
model.to(device)
base_lr = trial.suggest_uniform('base_lr', 1e-4, 1e-3)
lr_ratio = 10
ignored_params = list(map(id, model.classifier.parameters()))
print('the num of new layers:', len(ignored_params), flush=True)
base_params = filter(lambda p: id(p) not in ignored_params,
model.parameters())
optimizer = optim.SGD([{
'params': base_params
}, {
'params': model.classifier.parameters(),
'lr': lr_ratio * base_lr
}],
lr=base_lr,
momentum=0.9)
acc = train_model(model, criterion, optimizer)
return acc
def __init__(self,
device,
do_erasing=False,
do_padding=False,
side_thres=1.6,
save_probs=False,
prob_period=2,
weighted_prob=False,
conf_thres=0.8,
very_conf_thres=0.9,
output_test_image=False):
# url = "https://github.com/harry0412xd/Dashcam_anomaly_detection/releases/download/v1.0/gluon_seresnext101_32x4d-244_checkpoint-69.pth.tar"
# checkpoint_path = "model_data/gluon_seresnext101_32x4d-244_checkpoint-69.pth.tar"
url = "https://github.com/harry0412xd/Dashcam_anomaly_detection/releases/download/v1.0/gluon_seresnext101_32x4d-244_checkpoint-228_cleaned.pth-c36f9352.pth"
checkpoint_path = "model_data/gluon_seresnext101_32x4d-244_checkpoint-228_cleaned.pth-c36f9352.pth"
if not os.path.isfile(checkpoint_path):
torch.utils.model_zoo.load_url(url, model_dir="model_data/")
model = create_model('gluon_seresnext101_32x4d',
num_classes=2,
checkpoint_path=checkpoint_path)
model.to(device)
model.eval()
self.device = device
self.model = model
self.transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])
self.checkpoint_path = checkpoint_path
self.test_counter = {}
self.output_test_image = output_test_image
self.do_erasing = do_erasing
self.do_padding = do_padding
self.prob_period = prob_period # |<- n-th before --- current --- nth after->|
self.side_thres = side_thres
self.save_probs = save_probs
self.weighted_prob = weighted_prob
self.id2probs = {}
self.conf_thres = conf_thres
self.very_conf_thres = very_conf_thres
def model_creator(config):
args = config["args"]
model = create_model(
args.model,
pretrained=args.pretrained,
num_classes=args.num_classes,
drop_rate=args.drop,
drop_connect_rate=args.drop_connect, # DEPRECATED, use drop_path
drop_path_rate=args.drop_path,
drop_block_rate=args.drop_block,
global_pool=args.gp,
bn_tf=args.bn_tf,
bn_momentum=args.bn_momentum,
bn_eps=args.bn_eps,
checkpoint_path=args.initial_checkpoint)
# always false right now
if args.split_bn:
assert args.num_aug_splits > 1 or args.resplit
model = convert_splitbn_model(model, max(args.num_aug_splits, 2))
return model
def __init__(self,
save_path=None,
train_batch_size=256,
test_batch_size=512,
valid_size=None,
n_worker=32,
resize_scale=0.08,
distort_color=None,
image_size=224,
tf_preprocessing=False,
num_replicas=None,
rank=None,
use_prefetcher=False,
pin_memory=False,
fp16=False):
warnings.filterwarnings('ignore')
dataset = Dataset(os.path.join(save_path, "val"),
load_bytes=tf_preprocessing)
dummy_model = create_model('efficientnet_b0')
data_config = resolve_data_config({}, model=dummy_model)
test_loader = create_loader(dataset,
input_size=image_size,
batch_size=test_batch_size,
use_prefetcher=use_prefetcher,
interpolation=data_config['interpolation'],
mean=data_config['mean'],
std=data_config['std'],
num_workers=n_worker,
crop_pct=data_config['crop_pct'],
pin_memory=pin_memory,
fp16=fp16,
tf_preprocessing=None)
self.test = test_loader
def main():
args = parse_args()
if len(args.shape) == 1:
input_shape = (3, args.shape[0], args.shape[0])
elif len(args.shape) == 2:
input_shape = (3, ) + tuple(args.shape)
else:
raise ValueError('invalid input shape')
model = create_model(args.model, pretrained=False, num_classes=1000)
model.name = args.model
if torch.cuda.is_available():
model.cuda()
model.eval()
flops, params = get_flops(model, input_shape)
split_line = '=' * 30
print(f'{split_line}\nInput shape: {input_shape}\n'
f'Flops: {flops}\nParams: {params}\n{split_line}')
print('!!!Please be cautious if you use the results in papers. '
'You may need to check if all ops are supported and verify that the '
'flops computation is correct.')
def validate(args):
# might as well try to validate something
args.pretrained = args.pretrained or not args.checkpoint
args.prefetcher = not args.no_prefetcher
# create model
model = create_model(
args.model,
num_classes=args.num_classes,
in_chans=3,
pretrained=args.pretrained)
if args.checkpoint:
load_checkpoint(model, args.checkpoint, args.use_ema)
param_count = sum([m.numel() for m in model.parameters()])
logging.info('Model %s created, param count: %d' % (args.model, param_count))
data_config = resolve_data_config(vars(args), model=model)
model, test_time_pool = apply_test_time_pool(model, data_config, args)
if args.torchscript:
torch.jit.optimized_execution(True)
model = torch.jit.script(model)
if args.amp:
model = amp.initialize(model.cuda(), opt_level='O1')
else:
model = model.cuda()
if args.num_gpu > 1:
model = torch.nn.DataParallel(model, device_ids=list(range(args.num_gpu)))
criterion = nn.CrossEntropyLoss().cuda()
#from torchvision.datasets import ImageNet
#dataset = ImageNet(args.data, split='val')
if os.path.splitext(args.data)[1] == '.tar' and os.path.isfile(args.data):
dataset = DatasetTar(args.data, load_bytes=args.tf_preprocessing, class_map=args.class_map)
else:
dataset = Dataset(args.data, load_bytes=args.tf_preprocessing, class_map=args.class_map)
crop_pct = 1.0 if test_time_pool else data_config['crop_pct']
loader = create_loader(
dataset,
input_size=data_config['input_size'],
batch_size=args.batch_size,
use_prefetcher=args.prefetcher,
interpolation=data_config['interpolation'],
mean=data_config['mean'],
std=data_config['std'],
num_workers=args.workers,
crop_pct=crop_pct,
pin_memory=args.pin_mem,
tf_preprocessing=args.tf_preprocessing)
batch_time = AverageMeter()
losses = AverageMeter()
top1 = AverageMeter()
top5 = AverageMeter()
model.eval()
end = time.time()
with torch.no_grad():
for i, (input, target) in enumerate(loader):
if args.no_prefetcher:
target = target.cuda()
input = input.cuda()
if args.fp16:
input = input.half()
# compute output
output = model(input)
loss = criterion(output, target)
# measure accuracy and record loss
acc1, acc5 = accuracy(output.data, target, topk=(1, 2))
losses.update(loss.item(), input.size(0))
top1.update(acc1.item(), input.size(0))
top5.update(acc5.item(), input.size(0))
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if i % args.log_freq == 0:
logging.info(
'Test: [{0:>4d}/{1}] '
'Time: {batch_time.val:.3f}s ({batch_time.avg:.3f}s, {rate_avg:>7.2f}/s) '
'Loss: {loss.val:>7.4f} ({loss.avg:>6.4f}) '
'Acc@1: {top1.val:>7.3f} ({top1.avg:>7.3f}) '
'Acc@5: {top5.val:>7.3f} ({top5.avg:>7.3f})'.format(
i, len(loader), batch_time=batch_time,
rate_avg=input.size(0) / batch_time.avg,
loss=losses, top1=top1, top5=top5))
results = OrderedDict(
top1=round(top1.avg, 4), top1_err=round(100 - top1.avg, 4),
top5=round(top5.avg, 4), top5_err=round(100 - top5.avg, 4),
param_count=round(param_count / 1e6, 2),
img_size=data_config['input_size'][-1],
cropt_pct=crop_pct,
interpolation=data_config['interpolation'])
logging.info(' * Acc@1 {:.3f} ({:.3f}) Acc@5 {:.3f} ({:.3f})'.format(
results['top1'], results['top1_err'], results['top5'], results['top5_err']))
return results
def main():
setup_default_logging()
args, args_text = _parse_args()
args.prefetcher = not args.no_prefetcher
args.distributed = False
if 'WORLD_SIZE' in os.environ:
args.distributed = int(os.environ['WORLD_SIZE']) > 1
if args.distributed and args.num_gpu > 1:
_logger.warning(
'Using more than one GPU per process in distributed mode is not allowed.Setting num_gpu to 1.')
args.num_gpu = 1
args.device = 'cuda:0'
args.world_size = 1
args.rank = 0 # global rank
if args.distributed:
args.num_gpu = 1
args.device = 'cuda:%d' % args.local_rank
torch.cuda.set_device(args.local_rank)
torch.distributed.init_process_group(backend='nccl', init_method='env://')
args.world_size = torch.distributed.get_world_size()
args.rank = torch.distributed.get_rank()
assert args.rank >= 0
if args.distributed:
_logger.info('Training in distributed mode with multiple processes, 1 GPU per process. Process %d, total %d.'
% (args.rank, args.world_size))
else:
_logger.info('Training with a single process on %d GPUs.' % args.num_gpu)
torch.manual_seed(args.seed + args.rank)
model = create_model(
args.model,
pretrained=args.pretrained,
num_classes=args.num_classes,
drop_rate=args.drop,
drop_connect_rate=args.drop_connect, # DEPRECATED, use drop_path
drop_path_rate=args.drop_path,
drop_block_rate=args.drop_block,
global_pool=args.gp,
bn_tf=args.bn_tf,
bn_momentum=args.bn_momentum,
bn_eps=args.bn_eps,
checkpoint_path=args.initial_checkpoint)
if args.local_rank == 0:
_logger.info('Model %s created, param count: %d' %
(args.model, sum([m.numel() for m in model.parameters()])))
data_config = resolve_data_config(vars(args), model=model, verbose=args.local_rank == 0)
num_aug_splits = 0
if args.aug_splits > 0:
assert args.aug_splits > 1, 'A split of 1 makes no sense'
num_aug_splits = args.aug_splits
if args.split_bn:
assert num_aug_splits > 1 or args.resplit
model = convert_splitbn_model(model, max(num_aug_splits, 2))
use_amp = None
if args.amp:
# for backwards compat, `--amp` arg tries apex before native amp
if has_apex:
args.apex_amp = True
elif has_native_amp:
args.native_amp = True
if args.apex_amp and has_apex:
use_amp = 'apex'
elif args.native_amp and has_native_amp:
use_amp = 'native'
elif args.apex_amp or args.native_amp:
_logger.warning("Neither APEX or native Torch AMP is available, using float32. "
"Install NVIDA apex or upgrade to PyTorch 1.6")
if args.num_gpu > 1:
if use_amp == 'apex':
_logger.warning(
'Apex AMP does not work well with nn.DataParallel, disabling. Use DDP or Torch AMP.')
use_amp = None
model = nn.DataParallel(model, device_ids=list(range(args.num_gpu))).cuda()
assert not args.channels_last, "Channels last not supported with DP, use DDP."
else:
model.cuda()
if args.channels_last:
model = model.to(memory_format=torch.channels_last)
optimizer = create_optimizer(args, model)
amp_autocast = suppress # do nothing
loss_scaler = None
if use_amp == 'apex':
model, optimizer = amp.initialize(model, optimizer, opt_level='O1')
loss_scaler = ApexScaler()
if args.local_rank == 0:
_logger.info('Using NVIDIA APEX AMP. Training in mixed precision.')
elif use_amp == 'native':
amp_autocast = torch.cuda.amp.autocast
loss_scaler = NativeScaler()
if args.local_rank == 0:
_logger.info('Using native Torch AMP. Training in mixed precision.')
else:
if args.local_rank == 0:
_logger.info('AMP not enabled. Training in float32.')
# optionally resume from a checkpoint
resume_epoch = None
if args.resume:
resume_epoch = resume_checkpoint(
model, args.resume,
optimizer=None if args.no_resume_opt else optimizer,
loss_scaler=None if args.no_resume_opt else loss_scaler,
log_info=args.local_rank == 0)
model_ema = None
if args.model_ema:
# Important to create EMA model after cuda(), DP wrapper, and AMP but before SyncBN and DDP wrapper
model_ema = ModelEma(
model,
decay=args.model_ema_decay,
device='cpu' if args.model_ema_force_cpu else '',
resume=args.resume)
if args.distributed:
if args.sync_bn:
assert not args.split_bn
try:
if has_apex and use_amp != 'native':
# Apex SyncBN preferred unless native amp is activated
model = convert_syncbn_model(model)
else:
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
if args.local_rank == 0:
_logger.info(
'Converted model to use Synchronized BatchNorm. WARNING: You may have issues if using '
'zero initialized BN layers (enabled by default for ResNets) while sync-bn enabled.')
except Exception as e:
_logger.error('Failed to enable Synchronized BatchNorm. Install Apex or Torch >= 1.1')
if has_apex and use_amp != 'native':
# Apex DDP preferred unless native amp is activated
if args.local_rank == 0:
_logger.info("Using NVIDIA APEX DistributedDataParallel.")
model = ApexDDP(model, delay_allreduce=True)
else:
if args.local_rank == 0:
_logger.info("Using native Torch DistributedDataParallel.")
model = NativeDDP(model, device_ids=[args.local_rank]) # can use device str in Torch >= 1.1
# NOTE: EMA model does not need to be wrapped by DDP
lr_scheduler, num_epochs = create_scheduler(args, optimizer)
start_epoch = 0
if args.start_epoch is not None:
# a specified start_epoch will always override the resume epoch
start_epoch = args.start_epoch
elif resume_epoch is not None:
start_epoch = resume_epoch
if lr_scheduler is not None and start_epoch > 0:
lr_scheduler.step(start_epoch)
if args.local_rank == 0:
_logger.info('Scheduled epochs: {}'.format(num_epochs))
train_dir = os.path.join(args.data, 'train')
if not os.path.exists(train_dir):
_logger.error('Training folder does not exist at: {}'.format(train_dir))
exit(1)
dataset_train = Dataset(train_dir)
collate_fn = None
mixup_fn = None
mixup_active = args.mixup > 0 or args.cutmix > 0. or args.cutmix_minmax is not None
if mixup_active:
mixup_args = dict(
mixup_alpha=args.mixup, cutmix_alpha=args.cutmix, cutmix_minmax=args.cutmix_minmax,
prob=args.mixup_prob, switch_prob=args.mixup_switch_prob, mode=args.mixup_mode,
label_smoothing=args.smoothing, num_classes=args.num_classes)
if args.prefetcher:
assert not num_aug_splits # collate conflict (need to support deinterleaving in collate mixup)
collate_fn = FastCollateMixup(**mixup_args)
else:
mixup_fn = Mixup(**mixup_args)
if num_aug_splits > 1:
dataset_train = AugMixDataset(dataset_train, num_splits=num_aug_splits)
train_interpolation = args.train_interpolation
if args.no_aug or not train_interpolation:
train_interpolation = data_config['interpolation']
loader_train = create_loader(
dataset_train,
input_size=data_config['input_size'],
batch_size=args.batch_size,
is_training=True,
use_prefetcher=args.prefetcher,
no_aug=args.no_aug,
re_prob=args.reprob,
re_mode=args.remode,
re_count=args.recount,
re_split=args.resplit,
scale=args.scale,
ratio=args.ratio,
hflip=args.hflip,
vflip=args.vflip,
color_jitter=args.color_jitter,
auto_augment=args.aa,
num_aug_splits=num_aug_splits,
interpolation=train_interpolation,
mean=data_config['mean'],
std=data_config['std'],
num_workers=args.workers,
distributed=args.distributed,
collate_fn=collate_fn,
pin_memory=args.pin_mem,
use_multi_epochs_loader=args.use_multi_epochs_loader
)
eval_dir = os.path.join(args.data, 'val')
if not os.path.isdir(eval_dir):
eval_dir = os.path.join(args.data, 'validation')
if not os.path.isdir(eval_dir):
_logger.error('Validation folder does not exist at: {}'.format(eval_dir))
exit(1)
dataset_eval = Dataset(eval_dir)
loader_eval = create_loader(
dataset_eval,
input_size=data_config['input_size'],
batch_size=args.validation_batch_size_multiplier * args.batch_size,
is_training=False,
use_prefetcher=args.prefetcher,
interpolation=data_config['interpolation'],
mean=data_config['mean'],
std=data_config['std'],
num_workers=args.workers,
distributed=args.distributed,
crop_pct=data_config['crop_pct'],
pin_memory=args.pin_mem,
)
if args.jsd:
assert num_aug_splits > 1 # JSD only valid with aug splits set
train_loss_fn = JsdCrossEntropy(num_splits=num_aug_splits, smoothing=args.smoothing).cuda()
elif mixup_active:
# smoothing is handled with mixup target transform
train_loss_fn = SoftTargetCrossEntropy().cuda()
elif args.smoothing:
train_loss_fn = LabelSmoothingCrossEntropy(smoothing=args.smoothing).cuda()
else:
train_loss_fn = nn.CrossEntropyLoss().cuda()
validate_loss_fn = nn.CrossEntropyLoss().cuda()
eval_metric = args.eval_metric
best_metric = None
best_epoch = None
saver = None
output_dir = ''
if args.local_rank == 0:
output_base = args.output if args.output else './output'
exp_name = '-'.join([
datetime.now().strftime("%Y%m%d-%H%M%S"),
args.model,
str(data_config['input_size'][-1])
])
output_dir = get_outdir(output_base, 'train', exp_name)
decreasing = True if eval_metric == 'loss' else False
saver = CheckpointSaver(
model=model, optimizer=optimizer, args=args, model_ema=model_ema, amp_scaler=loss_scaler,
checkpoint_dir=output_dir, recovery_dir=output_dir, decreasing=decreasing)
with open(os.path.join(output_dir, 'args.yaml'), 'w') as f:
f.write(args_text)
try:
for epoch in range(start_epoch, num_epochs):
if args.distributed:
loader_train.sampler.set_epoch(epoch)
train_metrics = train_epoch(
epoch, model, loader_train, optimizer, train_loss_fn, args,
lr_scheduler=lr_scheduler, saver=saver, output_dir=output_dir,
amp_autocast=amp_autocast, loss_scaler=loss_scaler, model_ema=model_ema, mixup_fn=mixup_fn)
if args.distributed and args.dist_bn in ('broadcast', 'reduce'):
if args.local_rank == 0:
_logger.info("Distributing BatchNorm running means and vars")
distribute_bn(model, args.world_size, args.dist_bn == 'reduce')
eval_metrics = validate(model, loader_eval, validate_loss_fn, args, amp_autocast=amp_autocast)
if model_ema is not None and not args.model_ema_force_cpu:
if args.distributed and args.dist_bn in ('broadcast', 'reduce'):
distribute_bn(model_ema, args.world_size, args.dist_bn == 'reduce')
ema_eval_metrics = validate(
model_ema.ema, loader_eval, validate_loss_fn, args, amp_autocast=amp_autocast, log_suffix=' (EMA)')
eval_metrics = ema_eval_metrics
if lr_scheduler is not None:
# step LR for next epoch
lr_scheduler.step(epoch + 1, eval_metrics[eval_metric])
update_summary(
epoch, train_metrics, eval_metrics, os.path.join(output_dir, 'summary.csv'),
write_header=best_metric is None)
if saver is not None:
# save proper checkpoint with eval metric
save_metric = eval_metrics[eval_metric]
best_metric, best_epoch = saver.save_checkpoint(epoch, metric=save_metric)
# if saver.cmp(best_metric, save_metric):
# _logger.info(f"Metric is no longer improving [BEST: {best_metric}, CURRENT: {save_metric}]"
# f"\nFinishing training process")
# if epoch > 15:
# break
except KeyboardInterrupt:
pass
if best_metric is not None:
message = '*** Best metric: <{0:.2f}>, epoch: <{1}>, path: <{2}> ***'\
.format(best_metric, best_epoch, output_dir)
_logger.info(message)
print(message)
def main(args):
utils.init_distributed_mode(args)
print(args)
if args.distillation_type != 'none' and args.finetune and not args.eval:
raise NotImplementedError(
"Finetuning with distillation not yet supported")
device = torch.device(args.device)
# fix the seed for reproducibility
seed = args.seed + utils.get_rank()
torch.manual_seed(seed)
np.random.seed(seed)
# random.seed(seed)
cudnn.benchmark = True
dataset_train, args.nb_classes = build_dataset(is_train=True, args=args)
dataset_val, _ = build_dataset(is_train=False, args=args)
if True: # args.distributed:
num_tasks = utils.get_world_size()
global_rank = utils.get_rank()
if args.repeated_aug:
sampler_train = RASampler(dataset_train,
num_replicas=num_tasks,
rank=global_rank,
shuffle=True)
else:
sampler_train = torch.utils.data.DistributedSampler(
dataset_train,
num_replicas=num_tasks,
rank=global_rank,
shuffle=True)
if args.dist_eval:
if len(dataset_val) % num_tasks != 0:
print(
'Warning: Enabling distributed evaluation with an eval dataset not divisible by process number. '
'This will slightly alter validation results as extra duplicate entries are added to achieve '
'equal num of samples per-process.')
sampler_val = torch.utils.data.DistributedSampler(
dataset_val,
num_replicas=num_tasks,
rank=global_rank,
shuffle=False)
else:
sampler_val = torch.utils.data.SequentialSampler(dataset_val)
else:
sampler_train = torch.utils.data.RandomSampler(dataset_train)
sampler_val = torch.utils.data.SequentialSampler(dataset_val)
data_loader_train = torch.utils.data.DataLoader(
dataset_train,
sampler=sampler_train,
batch_size=args.batch_size,
num_workers=args.num_workers,
pin_memory=args.pin_mem,
drop_last=True,
)
data_loader_val = torch.utils.data.DataLoader(dataset_val,
sampler=sampler_val,
batch_size=int(
1.5 * args.batch_size),
num_workers=args.num_workers,
pin_memory=args.pin_mem,
drop_last=False)
mixup_fn = None
mixup_active = args.mixup > 0 or args.cutmix > 0. or args.cutmix_minmax is not None
if mixup_active:
mixup_fn = Mixup(mixup_alpha=args.mixup,
cutmix_alpha=args.cutmix,
cutmix_minmax=args.cutmix_minmax,
prob=args.mixup_prob,
switch_prob=args.mixup_switch_prob,
mode=args.mixup_mode,
label_smoothing=args.smoothing,
num_classes=args.nb_classes)
print(f"Creating model: {args.model}")
model = create_model(
args.model,
pretrained=False,
num_classes=args.nb_classes,
drop_rate=args.drop,
drop_path_rate=args.drop_path,
drop_block_rate=None,
)
if args.finetune:
if args.finetune.startswith('https'):
checkpoint = torch.hub.load_state_dict_from_url(args.finetune,
map_location='cpu',
check_hash=True)
else:
checkpoint = torch.load(args.finetune, map_location='cpu')
checkpoint_model = checkpoint['model']
state_dict = model.state_dict()
for k in [
'head.weight', 'head.bias', 'head_dist.weight',
'head_dist.bias'
]:
if k in checkpoint_model and checkpoint_model[
k].shape != state_dict[k].shape:
print(f"Removing key {k} from pretrained checkpoint")
del checkpoint_model[k]
# interpolate position embedding
pos_embed_checkpoint = checkpoint_model['pos_embed']
embedding_size = pos_embed_checkpoint.shape[-1]
num_patches = model.patch_embed.num_patches
num_extra_tokens = model.pos_embed.shape[-2] - num_patches
# height (== width) for the checkpoint position embedding
orig_size = int(
(pos_embed_checkpoint.shape[-2] - num_extra_tokens)**0.5)
# height (== width) for the new position embedding
new_size = int(num_patches**0.5)
# class_token and dist_token are kept unchanged
extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens]
# only the position tokens are interpolated
pos_tokens = pos_embed_checkpoint[:, num_extra_tokens:]
pos_tokens = pos_tokens.reshape(-1, orig_size, orig_size,
embedding_size).permute(0, 3, 1, 2)
pos_tokens = torch.nn.functional.interpolate(pos_tokens,
size=(new_size, new_size),
mode='bicubic',
align_corners=False)
pos_tokens = pos_tokens.permute(0, 2, 3, 1).flatten(1, 2)
new_pos_embed = torch.cat((extra_tokens, pos_tokens), dim=1)
checkpoint_model['pos_embed'] = new_pos_embed
model.load_state_dict(checkpoint_model, strict=False)
model.to(device)
model_ema = None
if args.model_ema:
# Important to create EMA model after cuda(), DP wrapper, and AMP but before SyncBN and DDP wrapper
model_ema = ModelEma(model,
decay=args.model_ema_decay,
device='cpu' if args.model_ema_force_cpu else '',
resume='')
model_without_ddp = model
if args.distributed:
model = torch.nn.parallel.DistributedDataParallel(
model, device_ids=[args.gpu])
model_without_ddp = model.module
n_parameters = sum(p.numel() for p in model.parameters()
if p.requires_grad)
print('number of params:', n_parameters)
linear_scaled_lr = args.lr * args.batch_size * utils.get_world_size(
) / 512.0
args.lr = linear_scaled_lr
optimizer = create_optimizer(args, model_without_ddp)
loss_scaler = NativeScaler()
lr_scheduler, _ = create_scheduler(args, optimizer)
criterion = LabelSmoothingCrossEntropy()
if args.mixup > 0.:
# smoothing is handled with mixup label transform
criterion = SoftTargetCrossEntropy()
elif args.smoothing:
criterion = LabelSmoothingCrossEntropy(smoothing=args.smoothing)
else:
criterion = torch.nn.CrossEntropyLoss()
teacher_model = None
if args.distillation_type != 'none':
assert args.teacher_path, 'need to specify teacher-path when using distillation'
print(f"Creating teacher model: {args.teacher_model}")
teacher_model = create_model(
args.teacher_model,
pretrained=False,
num_classes=args.nb_classes,
global_pool='avg',
)
if args.teacher_path.startswith('https'):
checkpoint = torch.hub.load_state_dict_from_url(args.teacher_path,
map_location='cpu',
check_hash=True)
else:
checkpoint = torch.load(args.teacher_path, map_location='cpu')
teacher_model.load_state_dict(checkpoint['model'])
teacher_model.to(device)
teacher_model.eval()
# wrap the criterion in our custom DistillationLoss, which
# just dispatches to the original criterion if args.distillation_type is 'none'
criterion = DistillationLoss(criterion, teacher_model,
args.distillation_type,
args.distillation_alpha,
args.distillation_tau)
output_dir = Path(args.output_dir)
if args.resume:
if args.resume.startswith('https'):
checkpoint = torch.hub.load_state_dict_from_url(args.resume,
map_location='cpu',
check_hash=True)
else:
checkpoint = torch.load(args.resume, map_location='cpu')
model_without_ddp.load_state_dict(checkpoint['model'])
if not args.eval and 'optimizer' in checkpoint and 'lr_scheduler' in checkpoint and 'epoch' in checkpoint:
optimizer.load_state_dict(checkpoint['optimizer'])
lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
args.start_epoch = checkpoint['epoch'] + 1
if args.model_ema:
utils._load_checkpoint_for_ema(model_ema,
checkpoint['model_ema'])
if 'scaler' in checkpoint:
loss_scaler.load_state_dict(checkpoint['scaler'])
if args.eval:
test_stats = evaluate(data_loader_val, model, device)
print(
f"Accuracy of the network on the {len(dataset_val)} test images: {test_stats['acc1']:.1f}%"
)
return
print(f"Start training for {args.epochs} epochs")
start_time = time.time()
max_accuracy = 0.0
for epoch in range(args.start_epoch, args.epochs):
if args.distributed:
data_loader_train.sampler.set_epoch(epoch)
train_stats = train_one_epoch(
model,
criterion,
data_loader_train,
optimizer,
device,
epoch,
loss_scaler,
args.clip_grad,
model_ema,
mixup_fn,
set_training_mode=args.finetune ==
'' # keep in eval mode during finetuning
)
lr_scheduler.step(epoch)
if args.output_dir:
checkpoint_paths = [output_dir / ('checkpoint_%04d.pth' % (epoch))]
for checkpoint_path in checkpoint_paths:
utils.save_on_master(
{
'model': model_without_ddp.state_dict(),
'optimizer': optimizer.state_dict(),
'lr_scheduler': lr_scheduler.state_dict(),
'epoch': epoch,
'model_ema': get_state_dict(model_ema),
'scaler': loss_scaler.state_dict(),
'args': args,
}, checkpoint_path)
if not args.train_without_eval:
test_stats = evaluate(data_loader_val, model, device)
print(
f"Accuracy of the network on the {len(dataset_val)} test images: {test_stats['acc1']:.1f}%"
)
max_accuracy = max(max_accuracy, test_stats["acc1"])
print(f'Max accuracy: {max_accuracy:.2f}%')
log_stats = {
**{f'train_{k}': v
for k, v in train_stats.items()},
**{f'test_{k}': v
for k, v in test_stats.items()}, 'epoch': epoch,
'n_parameters': n_parameters
}
else:
log_stats = {
**{f'train_{k}': v
for k, v in train_stats.items()}, 'epoch': epoch,
'n_parameters': n_parameters
}
if args.output_dir and utils.is_main_process():
with (output_dir / "log.txt").open("a") as f:
f.write(json.dumps(log_stats) + "\n")
total_time = time.time() - start_time
total_time_str = str(datetime.timedelta(seconds=int(total_time)))
print('Training time {}'.format(total_time_str))
def main():
setup_default_logging()
args, args_text = _parse_args()
if args.log_wandb:
if has_wandb:
wandb.init(project=args.experiment, config=args)
else:
_logger.warning(
"You've requested to log metrics to wandb but package not found. "
"Metrics not being logged to wandb, try `pip install wandb`")
args.prefetcher = not args.no_prefetcher
args.distributed = False
if 'WORLD_SIZE' in os.environ:
args.distributed = int(os.environ['WORLD_SIZE']) > 1
args.device = 'cuda:0'
args.world_size = 1
args.rank = 0 # global rank
if args.distributed:
args.device = 'cuda:%d' % args.local_rank
torch.cuda.set_device(args.local_rank)
torch.distributed.init_process_group(backend='nccl',
init_method='env://')
args.world_size = torch.distributed.get_world_size()
args.rank = torch.distributed.get_rank()
_logger.info(
'Training in distributed mode with multiple processes, 1 GPU per process. Process %d, total %d.'
% (args.rank, args.world_size))
else:
_logger.info('Training with a single process on 1 GPUs.')
assert args.rank >= 0
# resolve AMP arguments based on PyTorch / Apex availability
use_amp = None
if args.amp:
# `--amp` chooses native amp before apex (APEX ver not actively maintained)
if has_native_amp:
args.native_amp = True
elif has_apex:
args.apex_amp = True
if args.apex_amp and has_apex:
use_amp = 'apex'
elif args.native_amp and has_native_amp:
use_amp = 'native'
elif args.apex_amp or args.native_amp:
_logger.warning(
"Neither APEX or native Torch AMP is available, using float32. "
"Install NVIDA apex or upgrade to PyTorch 1.6")
random_seed(args.seed, args.rank)
model_KD = None
if args.kd_model_path is not None:
model_KD = build_kd_model(args)
model = create_model(
args.model,
pretrained=args.pretrained,
num_classes=args.num_classes,
drop_rate=args.drop,
drop_connect_rate=args.drop_connect, # DEPRECATED, use drop_path
drop_path_rate=args.drop_path,
drop_block_rate=args.drop_block,
global_pool=args.gp,
bn_momentum=args.bn_momentum,
bn_eps=args.bn_eps,
scriptable=args.torchscript,
checkpoint_path=args.initial_checkpoint)
if args.num_classes is None:
assert hasattr(
model, 'num_classes'
), 'Model must have `num_classes` attr if not set on cmd line/config.'
args.num_classes = model.num_classes # FIXME handle model default vs config num_classes more elegantly
if args.local_rank == 0:
_logger.info(
f'Model {safe_model_name(args.model)} created, param count:{sum([m.numel() for m in model.parameters()])}'
)
data_config = resolve_data_config(vars(args),
model=model,
verbose=args.local_rank == 0)
# setup augmentation batch splits for contrastive loss or split bn
num_aug_splits = 0
if args.aug_splits > 0:
assert args.aug_splits > 1, 'A split of 1 makes no sense'
num_aug_splits = args.aug_splits
# enable split bn (separate bn stats per batch-portion)
if args.split_bn:
assert num_aug_splits > 1 or args.resplit
model = convert_splitbn_model(model, max(num_aug_splits, 2))
# move model to GPU, enable channels last layout if set
model.cuda()
if args.channels_last:
model = model.to(memory_format=torch.channels_last)
# setup synchronized BatchNorm for distributed training
if args.distributed and args.sync_bn:
assert not args.split_bn
if has_apex and use_amp == 'apex':
# Apex SyncBN preferred unless native amp is activated
model = convert_syncbn_model(model)
else:
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
if args.local_rank == 0:
_logger.info(
'Converted model to use Synchronized BatchNorm. WARNING: You may have issues if using '
'zero initialized BN layers (enabled by default for ResNets) while sync-bn enabled.'
)
if args.torchscript:
assert not use_amp == 'apex', 'Cannot use APEX AMP with torchscripted model'
assert not args.sync_bn, 'Cannot use SyncBatchNorm with torchscripted model'
model = torch.jit.script(model)
optimizer = create_optimizer_v2(model, **optimizer_kwargs(cfg=args))
# setup automatic mixed-precision (AMP) loss scaling and op casting
amp_autocast = suppress # do nothing
loss_scaler = None
if use_amp == 'apex':
model, optimizer = amp.initialize(model, optimizer, opt_level='O1')
loss_scaler = ApexScaler()
if args.local_rank == 0:
_logger.info('Using NVIDIA APEX AMP. Training in mixed precision.')
elif use_amp == 'native':
amp_autocast = torch.cuda.amp.autocast
loss_scaler = NativeScaler()
if args.local_rank == 0:
_logger.info(
'Using native Torch AMP. Training in mixed precision.')
else:
if args.local_rank == 0:
_logger.info('AMP not enabled. Training in float32.')
# optionally resume from a checkpoint
resume_epoch = None
if args.resume:
resume_epoch = resume_checkpoint(
model,
args.resume,
optimizer=None if args.no_resume_opt else optimizer,
loss_scaler=None if args.no_resume_opt else loss_scaler,
log_info=args.local_rank == 0)
# setup exponential moving average of model weights, SWA could be used here too
model_ema = None
if args.model_ema:
# Important to create EMA model after cuda(), DP wrapper, and AMP but before SyncBN and DDP wrapper
model_ema = ModelEmaV2(
model,
decay=args.model_ema_decay,
device='cpu' if args.model_ema_force_cpu else None)
if args.resume:
load_checkpoint(model_ema.module, args.resume, use_ema=True)
# setup distributed training
if args.distributed:
if has_apex and use_amp == 'apex':
# Apex DDP preferred unless native amp is activated
if args.local_rank == 0:
_logger.info("Using NVIDIA APEX DistributedDataParallel.")
model = ApexDDP(model, delay_allreduce=True)
else:
if args.local_rank == 0:
_logger.info("Using native Torch DistributedDataParallel.")
model = NativeDDP(model,
device_ids=[args.local_rank],
broadcast_buffers=not args.no_ddp_bb)
# NOTE: EMA model does not need to be wrapped by DDP
# setup learning rate schedule and starting epoch
lr_scheduler, num_epochs = create_scheduler(args, optimizer)
start_epoch = 0
if args.start_epoch is not None:
# a specified start_epoch will always override the resume epoch
start_epoch = args.start_epoch
elif resume_epoch is not None:
start_epoch = resume_epoch
if lr_scheduler is not None and start_epoch > 0:
lr_scheduler.step(start_epoch)
if args.local_rank == 0:
_logger.info('Scheduled epochs: {}'.format(num_epochs))
# create the train and eval datasets
dataset_train = create_dataset(args.dataset,
root=args.data_dir,
split=args.train_split,
is_training=True,
class_map=args.class_map,
download=args.dataset_download,
batch_size=args.batch_size,
repeats=args.epoch_repeats)
dataset_eval = create_dataset(args.dataset,
root=args.data_dir,
split=args.val_split,
is_training=False,
class_map=args.class_map,
download=args.dataset_download,
batch_size=args.batch_size)
# setup mixup / cutmix
collate_fn = None
mixup_fn = None
mixup_active = args.mixup > 0 or args.cutmix > 0. or args.cutmix_minmax is not None
if mixup_active:
mixup_args = dict(mixup_alpha=args.mixup,
cutmix_alpha=args.cutmix,
cutmix_minmax=args.cutmix_minmax,
prob=args.mixup_prob,
switch_prob=args.mixup_switch_prob,
mode=args.mixup_mode,
label_smoothing=args.smoothing,
num_classes=args.num_classes)
if args.prefetcher:
assert not num_aug_splits # collate conflict (need to support deinterleaving in collate mixup)
collate_fn = FastCollateMixup(**mixup_args)
else:
mixup_fn = Mixup(**mixup_args)
# wrap dataset in AugMix helper
if num_aug_splits > 1:
dataset_train = AugMixDataset(dataset_train, num_splits=num_aug_splits)
# create data loaders w/ augmentation pipeiine
train_interpolation = args.train_interpolation
if args.no_aug or not train_interpolation:
train_interpolation = data_config['interpolation']
loader_train = create_loader(
dataset_train,
input_size=data_config['input_size'],
batch_size=args.batch_size,
is_training=True,
use_prefetcher=args.prefetcher,
no_aug=args.no_aug,
re_prob=args.reprob,
re_mode=args.remode,
re_count=args.recount,
re_split=args.resplit,
scale=args.scale,
ratio=args.ratio,
hflip=args.hflip,
vflip=args.vflip,
color_jitter=args.color_jitter,
auto_augment=args.aa,
num_aug_repeats=args.aug_repeats,
num_aug_splits=num_aug_splits,
interpolation=train_interpolation,
mean=data_config['mean'],
std=data_config['std'],
num_workers=args.workers,
distributed=args.distributed,
collate_fn=collate_fn,
pin_memory=args.pin_mem,
use_multi_epochs_loader=args.use_multi_epochs_loader,
worker_seeding=args.worker_seeding,
)
loader_eval = create_loader(
dataset_eval,
input_size=data_config['input_size'],
batch_size=args.validation_batch_size or args.batch_size,
is_training=False,
use_prefetcher=args.prefetcher,
interpolation=data_config['interpolation'],
mean=data_config['mean'],
std=data_config['std'],
num_workers=args.workers,
distributed=args.distributed,
crop_pct=data_config['crop_pct'],
pin_memory=args.pin_mem,
)
# setup loss function
if args.jsd_loss:
assert num_aug_splits > 1 # JSD only valid with aug splits set
train_loss_fn = JsdCrossEntropy(num_splits=num_aug_splits,
smoothing=args.smoothing)
elif mixup_active:
# smoothing is handled with mixup target transform which outputs sparse, soft targets
if args.bce_loss:
train_loss_fn = BinaryCrossEntropy(
target_threshold=args.bce_target_thresh)
else:
train_loss_fn = SoftTargetCrossEntropy()
elif args.smoothing:
if args.bce_loss:
train_loss_fn = BinaryCrossEntropy(
smoothing=args.smoothing,
target_threshold=args.bce_target_thresh)
else:
train_loss_fn = LabelSmoothingCrossEntropy(
smoothing=args.smoothing)
else:
train_loss_fn = nn.CrossEntropyLoss()
train_loss_fn = train_loss_fn.cuda()
validate_loss_fn = nn.CrossEntropyLoss().cuda()
# setup checkpoint saver and eval metric tracking
eval_metric = args.eval_metric
best_metric = None
best_epoch = None
saver = None
output_dir = None
if args.rank == 0:
if args.experiment:
exp_name = args.experiment
else:
exp_name = '-'.join([
datetime.now().strftime("%Y%m%d-%H%M%S"),
safe_model_name(args.model),
str(data_config['input_size'][-1])
])
output_dir = get_outdir(
args.output if args.output else './output/train', exp_name)
decreasing = True if eval_metric == 'loss' else False
saver = CheckpointSaver(model=model,
optimizer=optimizer,
args=args,
model_ema=model_ema,
amp_scaler=loss_scaler,
checkpoint_dir=output_dir,
recovery_dir=output_dir,
decreasing=decreasing,
max_history=args.checkpoint_hist)
with open(os.path.join(output_dir, 'args.yaml'), 'w') as f:
f.write(args_text)
try:
for epoch in range(start_epoch, num_epochs):
if args.distributed and hasattr(loader_train.sampler, 'set_epoch'):
loader_train.sampler.set_epoch(epoch)
train_metrics = train_one_epoch(epoch,
model,
loader_train,
optimizer,
train_loss_fn,
args,
lr_scheduler=lr_scheduler,
saver=saver,
output_dir=output_dir,
amp_autocast=amp_autocast,
loss_scaler=loss_scaler,
model_ema=model_ema,
mixup_fn=mixup_fn,
model_KD=model_KD)
if args.distributed and args.dist_bn in ('broadcast', 'reduce'):
if args.local_rank == 0:
_logger.info(
"Distributing BatchNorm running means and vars")
distribute_bn(model, args.world_size, args.dist_bn == 'reduce')
eval_metrics = validate(model,
loader_eval,
validate_loss_fn,
args,
amp_autocast=amp_autocast)
if model_ema is not None and not args.model_ema_force_cpu:
if args.distributed and args.dist_bn in ('broadcast',
'reduce'):
distribute_bn(model_ema, args.world_size,
args.dist_bn == 'reduce')
ema_eval_metrics = validate(model_ema.module,
loader_eval,
validate_loss_fn,
args,
amp_autocast=amp_autocast,
log_suffix=' (EMA)')
eval_metrics = ema_eval_metrics
if lr_scheduler is not None:
# step LR for next epoch
lr_scheduler.step(epoch + 1, eval_metrics[eval_metric])
if output_dir is not None:
update_summary(epoch,
train_metrics,
eval_metrics,
os.path.join(output_dir, 'summary.csv'),
write_header=best_metric is None,
log_wandb=args.log_wandb and has_wandb)
if saver is not None:
# save proper checkpoint with eval metric
save_metric = eval_metrics[eval_metric]
best_metric, best_epoch = saver.save_checkpoint(
epoch, metric=save_metric)
except KeyboardInterrupt:
pass
if best_metric is not None:
_logger.info('*** Best metric: {0} (epoch {1})'.format(
best_metric, best_epoch))
def validate(args):
# might as well try to validate something
args.pretrained = args.pretrained or not args.checkpoint
# create model
model = create_model(args.model,
pretrained=args.pretrained,
num_classes=args.num_classes,
in_chans=3,
scriptable=args.torchscript)
if args.checkpoint:
load_checkpoint(model, args.checkpoint, args.use_ema)
param_count = sum([m.numel() for m in model.parameters()])
logging.info('Model %s created, param count: %d' %
(args.model, param_count))
if args.torchscript:
torch.jit.optimized_execution(True)
model = torch.jit.script(model)
model = model.cuda()
if args.num_gpu > 1:
model = torch.nn.DataParallel(model,
device_ids=list(range(args.num_gpu)))
criterion = nn.CrossEntropyLoss().cuda()
# from torchvision.datasets import ImageNet
# dataset = ImageNet(args.data, split='val')
valdir = args.data
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
transform = cvtransforms.Compose([
cvtransforms.Resize(size=(256), interpolation='BILINEAR'),
cvtransforms.CenterCrop(224),
cvtransforms.ToTensor(),
cvtransforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])
# loader = torch.utils.data.DataLoader(
# datasets.ImageFolder(valdir, transform, loader=opencv_loader),
# batch_size=args.batch_size, shuffle=False,
# num_workers=args.workers, pin_memory=False)
loader = torch.utils.data.DataLoader(datasets.ImageFolder(
valdir,
transforms.Compose([
transforms.Resize((256), interpolation=2),
transforms.CenterCrop(224),
transforms.ToTensor(),
normalize,
])),
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=False)
# loader_eval = loader.Loader('val', valdir, batch_size=args.batch_size, num_workers=args.workers, shuffle=False)
batch_time = AverageMeter()
losses = AverageMeter()
top1 = AverageMeter()
top5 = AverageMeter()
model.eval()
with torch.no_grad():
# warmup, reduce variability of first batch time, especially for comparing torchscript vs non
# input = torch.randn((args.batch_size,)).cuda()
# model(input)
end = time.time()
for i, (input, target) in enumerate(loader):
# if args.no_prefetcher:
target = target.cuda()
input = input.cuda()
# compute output
output, _ = model(input)
# loss = criterion(output, target)
# measure accuracy and record loss
acc1, acc5 = accuracy(output.data, target, topk=(1, 5))
# losses.update(loss.item(), input.size(0))
top1.update(acc1.item(), input.size(0))
top5.update(acc5.item(), input.size(0))
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if i % args.log_freq == 0:
logging.info(
'Test: [{0:>4d}/{1}] '
'Time: {batch_time.val:.3f}s ({batch_time.avg:.3f}s, {rate_avg:>7.2f}/s) '
'Loss: {loss.val:>7.4f} ({loss.avg:>6.4f}) '
'Acc@1: {top1.val:>7.3f} ({top1.avg:>7.3f}) '
'Acc@5: {top5.val:>7.3f} ({top5.avg:>7.3f})'.format(
i,
len(loader),
batch_time=batch_time,
rate_avg=input.size(0) / batch_time.avg,
loss=losses,
top1=top1,
top5=top5))
results = OrderedDict(top1=round(top1.avg, 4),
top1_err=round(100 - top1.avg, 4),
top5=round(top5.avg, 4),
top5_err=round(100 - top5.avg, 4),
param_count=round(param_count / 1e6, 2))
logging.info(' * Acc@1 {:.3f} ({:.3f}) Acc@5 {:.3f} ({:.3f})'.format(
results['top1'], results['top1_err'], results['top5'],
results['top5_err']))
return results
def validate(args):
# might as well try to validate something
args.pretrained = args.pretrained or not args.checkpoint
args.prefetcher = not args.no_prefetcher
amp_autocast = suppress # do nothing
if args.amp:
if has_apex:
args.apex_amp = True
elif has_native_amp:
args.native_amp = True
else:
_logger.warning(
"Neither APEX or Native Torch AMP is available, using FP32.")
assert not args.apex_amp or not args.native_amp, "Only one AMP mode should be set."
if args.native_amp:
amp_autocast = torch.cuda.amp.autocast
if args.legacy_jit:
set_jit_legacy()
# create model
model = create_model(args.model,
pretrained=args.pretrained,
num_classes=args.num_classes,
in_chans=3,
global_pool=args.gp,
scriptable=args.torchscript)
if args.num_classes is None:
assert hasattr(
model, 'num_classes'
), 'Model must have `num_classes` attr if not set on cmd line/config.'
args.num_classes = model.num_classes
if args.checkpoint:
load_checkpoint(model, args.checkpoint, args.use_ema)
param_count = sum([m.numel() for m in model.parameters()])
_logger.info('Model %s created, param count: %d' %
(args.model, param_count))
data_config = resolve_data_config(vars(args), model=model)
model, test_time_pool = (
model, False) if args.no_test_pool else apply_test_time_pool(
model, data_config)
if args.torchscript:
torch.jit.optimized_execution(True)
model = torch.jit.script(model)
model = model.cuda()
if args.apex_amp:
model = amp.initialize(model, opt_level='O1')
if args.channels_last:
model = model.to(memory_format=torch.channels_last)
if args.num_gpu > 1:
model = torch.nn.DataParallel(model,
device_ids=list(range(args.num_gpu)))
criterion = nn.CrossEntropyLoss().cuda()
dataset = create_dataset(root=args.data,
name=args.dataset,
split=args.split,
load_bytes=args.tf_preprocessing,
class_map=args.class_map)
if args.valid_labels:
with open(args.valid_labels, 'r') as f:
valid_labels = {int(line.rstrip()) for line in f}
valid_labels = [i in valid_labels for i in range(args.num_classes)]
else:
valid_labels = None
if args.real_labels:
real_labels = RealLabelsImagenet(dataset.filenames(basename=True),
real_json=args.real_labels)
else:
real_labels = None
crop_pct = 1.0 if test_time_pool else data_config['crop_pct']
loader = create_loader(dataset,
input_size=data_config['input_size'],
batch_size=args.batch_size,
use_prefetcher=args.prefetcher,
interpolation=data_config['interpolation'],
mean=data_config['mean'],
std=data_config['std'],
num_workers=args.workers,
crop_pct=crop_pct,
pin_memory=args.pin_mem,
tf_preprocessing=args.tf_preprocessing)
batch_time = AverageMeter()
losses = AverageMeter()
top1 = AverageMeter()
top5 = AverageMeter()
model.eval()
with torch.no_grad():
# warmup, reduce variability of first batch time, especially for comparing torchscript vs non
input = torch.randn((args.batch_size, ) +
data_config['input_size']).cuda()
if args.channels_last:
input = input.contiguous(memory_format=torch.channels_last)
model(input)
end = time.time()
for batch_idx, (input, target) in enumerate(loader):
if args.no_prefetcher:
target = target.cuda()
input = input.cuda()
if args.channels_last:
input = input.contiguous(memory_format=torch.channels_last)
# compute output
with amp_autocast():
output = model(input)
if valid_labels is not None:
output = output[:, valid_labels]
loss = criterion(output, target)
if real_labels is not None:
real_labels.add_result(output)
# measure accuracy and record loss
acc1, acc5 = accuracy(output.detach(), target, topk=(1, 5))
losses.update(loss.item(), input.size(0))
top1.update(acc1.item(), input.size(0))
top5.update(acc5.item(), input.size(0))
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if batch_idx % args.log_freq == 0:
_logger.info(
'Test: [{0:>4d}/{1}] '
'Time: {batch_time.val:.3f}s ({batch_time.avg:.3f}s, {rate_avg:>7.2f}/s) '
'Loss: {loss.val:>7.4f} ({loss.avg:>6.4f}) '
'Acc@1: {top1.val:>7.3f} ({top1.avg:>7.3f}) '
'Acc@5: {top5.val:>7.3f} ({top5.avg:>7.3f})'.format(
batch_idx,
len(loader),
batch_time=batch_time,
rate_avg=input.size(0) / batch_time.avg,
loss=losses,
top1=top1,
top5=top5))
if real_labels is not None:
# real labels mode replaces topk values at the end
top1a, top5a = real_labels.get_accuracy(k=1), real_labels.get_accuracy(
k=5)
else:
top1a, top5a = top1.avg, top5.avg
results = OrderedDict(top1=round(top1a, 4),
top1_err=round(100 - top1a, 4),
top5=round(top5a, 4),
top5_err=round(100 - top5a, 4),
param_count=round(param_count / 1e6, 2),
img_size=data_config['input_size'][-1],
cropt_pct=crop_pct,
interpolation=data_config['interpolation'])
_logger.info(' * Acc@1 {:.3f} ({:.3f}) Acc@5 {:.3f} ({:.3f})'.format(
results['top1'], results['top1_err'], results['top5'],
results['top5_err']))
return results
def validate(args):
# might as well try to validate something
args.pretrained = args.pretrained or not args.checkpoint
args.prefetcher = not args.no_prefetcher
if args.legacy_jit:
set_jit_legacy()
# create model
if 'inception' in args.model:
model = create_model(
args.model,
pretrained=args.pretrained,
num_classes=args.num_classes,
aux_logits=True, # ! add aux loss
in_chans=3,
scriptable=args.torchscript)
else:
model = create_model(args.model,
pretrained=args.pretrained,
num_classes=args.num_classes,
in_chans=3,
scriptable=args.torchscript)
# ! add more layer to classifier layer
if args.create_classifier_layerfc:
model.global_pool, model.classifier = create_classifier_layerfc(
model.num_features, model.num_classes)
if args.checkpoint:
load_checkpoint(model, args.checkpoint, args.use_ema)
param_count = sum([m.numel() for m in model.parameters()])
_logger.info('Model %s created, param count: %d' %
(args.model, param_count))
data_config = resolve_data_config(vars(args), model=model)
model, test_time_pool = apply_test_time_pool(model, data_config, args)
if args.torchscript:
torch.jit.optimized_execution(True)
model = torch.jit.script(model)
if args.amp:
model = amp.initialize(model.cuda(), opt_level='O1')
else:
model = model.cuda()
if args.num_gpu > 1:
model = torch.nn.DataParallel(model,
device_ids=list(range(args.num_gpu)))
if args.has_eval_label:
criterion = nn.CrossEntropyLoss().cuda() # ! don't have gold label
if os.path.splitext(args.data)[1] == '.tar' and os.path.isfile(args.data):
dataset = DatasetTar(args.data,
load_bytes=args.tf_preprocessing,
class_map=args.class_map)
else:
dataset = Dataset(args.data,
load_bytes=args.tf_preprocessing,
class_map=args.class_map,
args=args)
if args.valid_labels:
with open(args.valid_labels,
'r') as f: # @valid_labels is index numbering
valid_labels = {int(line.rstrip()) for line in f}
valid_labels = [i in valid_labels for i in range(args.num_classes)]
else:
valid_labels = None
if args.real_labels:
real_labels = RealLabelsImagenet(dataset.filenames(basename=True),
real_json=args.real_labels)
else:
real_labels = None
crop_pct = 1.0 if test_time_pool else data_config['crop_pct']
loader = create_loader(
dataset,
input_size=data_config['input_size'],
batch_size=args.batch_size,
use_prefetcher=args.prefetcher,
interpolation=data_config[
'interpolation'], # 'blank' is default Image.BILINEAR https://github.com/rwightman/pytorch-image-models/blob/470220b1f4c61ad7deb16dbfb8917089e842cd2a/timm/data/transforms.py#L43
mean=data_config['mean'],
std=data_config['std'],
num_workers=args.workers,
crop_pct=crop_pct,
pin_memory=args.pin_mem,
tf_preprocessing=args.tf_preprocessing,
auto_augment=args.aa,
scale=args.scale,
ratio=args.ratio,
hflip=args.hflip,
vflip=args.vflip,
color_jitter=args.color_jitter,
args=args)
batch_time = AverageMeter()
losses = AverageMeter()
top1 = AverageMeter()
topk = AverageMeter()
prediction = None # ! need to save output
true_label = None
model.eval()
with torch.no_grad():
# warmup, reduce variability of first batch time, especially for comparing torchscript vs non
input = torch.randn((args.batch_size, ) +
data_config['input_size']).cuda()
model(input)
end = time.time()
for batch_idx, (input,
target) in enumerate(loader): # ! not have real label
if args.has_eval_label: # ! just save true labels anyway... why not
if true_label is None: true_label = target.cpu().data.numpy()
else:
true_label = np.concatenate(
(true_label, target.cpu().data.numpy()), axis=0)
if args.no_prefetcher:
target = target.cuda()
input = input.cuda()
if args.fp16:
input = input.half()
# compute output
output = model(input)
if isinstance(output, (tuple, list)):
output = output[0] # ! some model returns both loss + aux loss
if valid_labels is not None:
output = output[:,
valid_labels] # ! keep only valid labels ? good to eval by class.
# ! save prediction, don't append too slow ... whatever ?
# ! are names of files also sorted ?
if prediction is None:
prediction = output.cpu().data.numpy() # batchsize x label
else: # stack
prediction = np.concatenate(
(prediction, output.cpu().data.numpy()), axis=0)
if real_labels is not None:
real_labels.add_result(output)
if args.has_eval_label:
# measure accuracy and record loss
loss = criterion(
output, target) # ! don't have gold standard on testset
acc1, acc5 = accuracy(output.data, target, topk=(1, args.topk))
losses.update(loss.item(), input.size(0))
top1.update(acc1.item(), input.size(0))
topk.update(acc5.item(), input.size(0))
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if args.has_eval_label and (batch_idx % args.log_freq == 0):
_logger.info(
'Test: [{0:>4d}/{1}] '
'Time: {batch_time.val:.3f}s ({batch_time.avg:.3f}s, {rate_avg:>7.2f}/s) '
'Loss: {loss.val:>7.4f} ({loss.avg:>6.4f}) '
'Acc@1: {top1.val:>7.3f} ({top1.avg:>7.3f}) '
'Acc@topk: {topk.val:>7.3f} ({topk.avg:>7.3f})'.format(
batch_idx,
len(loader),
batch_time=batch_time,
rate_avg=input.size(0) / batch_time.avg,
loss=losses,
top1=top1,
topk=topk))
if not args.has_eval_label:
top1a, topka = 0, 0 # just dummy, because we don't know ground labels
else:
if real_labels is not None:
# real labels mode replaces topk values at the end
top1a, topka = real_labels.get_accuracy(
k=1), real_labels.get_accuracy(k=args.topk)
else:
top1a, topka = top1.avg, topk.avg
results = OrderedDict(top1=round(top1a, 4),
top1_err=round(100 - top1a, 4),
topk=round(topka, 4),
topk_err=round(100 - topka, 4),
param_count=round(param_count / 1e6, 2),
img_size=data_config['input_size'][-1],
cropt_pct=crop_pct,
interpolation=data_config['interpolation'])
_logger.info(' * Acc@1 {:.3f} ({:.3f}) Acc@topk {:.3f} ({:.3f})'.format(
results['top1'], results['top1_err'], results['topk'],
results['topk_err']))
return results, prediction, true_label
def __init__(
self,
epoch,
dataset_path='./drive/My Drive/datasets/car classification/train_dataset',
val_path='./drive/My Drive/datasets/car classification/val_data',
batch_size=128,
model_name='tf_efficientnet_b0_ns',
ckpt_path='./drive/My Drive/ckpt/190.pth',
test_number=5000,
pseudo_test=True,
crop='five',
csv_path='',
mode='fix',
sizes=(680, 600, 528)):
self.epoch = epoch
self.dataset_path = dataset_path
self.val_path = val_path
self.batch_size = batch_size
self.model_name = model_name
self.ckpt_path = ckpt_path
self.test_number = test_number
self.pseudo_test = pseudo_test
self.crop = crop
self.csv_path = csv_path
self.mode = mode
self.sizes = sizes
if model_name == 'tf_efficientnet_b0_ns':
self.input_size = (224, 224)
elif model_name == 'tf_efficientnet_b3_ns':
self.input_size = (300, 300)
elif model_name == 'tf_efficientnet_b4_ns':
self.input_size = (480, 480)
elif model_name == 'tf_efficientnet_b6_ns':
self.input_size = (680, 680) # 528
else:
raise Exception('non-valid model name')
# Compose transforms
transform = []
fill = lambda i: transforms.Resize((i.size[1] * (2**torch.ceil(
torch.log2(torch.tensor(self.input_size[1] / i.size[1]))
)), i.size[0] * (2**torch.ceil(
torch.log2(torch.tensor(self.input_size[1] / i.size[1]))))))(
i) if i.size[0] < self.input_size[0] or i.size[
1] < self.input_size[1] else i
if crop == 'center':
transform.append(transforms.CenterCrop(self.input_size[0]))
transform.append(transforms.ToTensor())
transform.append(
transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]))
elif crop == 'five':
transform.append(transforms.Lambda(fill))
transform.append(transforms.FiveCrop(self.input_size[0]))
transform.append(
transforms.Lambda(lambda crops: torch.stack(
[transforms.ToTensor()(crop) for crop in crops])))
transform.append(
transforms.Lambda(lambda crops: torch.stack([
transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
(crop) for crop in crops
])))
self.transform = transforms.Compose(transform)
if self.pseudo_test:
if crop == 'multi':
self.transform_val = []
self.dataset = []
self.dataloader = []
for i in range(len(self.sizes)):
self.transform_val.append(
self.get_transform_val((self.sizes[i], self.sizes[i])))
self.dataset.append(
ImageFolder(self.dataset_path,
transform=self.transform_val[i]))
self.dataloader.append(
DataLoader(self.dataset[i],
batch_size=self.batch_size,
num_workers=1,
shuffle=False))
else:
self.dataset = ImageFolder(self.dataset_path,
transform=self.transform_val)
self.dataloader = DataLoader(self.dataset,
batch_size=self.batch_size,
num_workers=1,
shuffle=False)
self.device = torch.device(
'cuda' if torch.cuda.is_available() else 'cpu')
self.model = create_model(model_name, num_classes=196).to(self.device)
if self.mode == 'fix':
ckpt = torch.load(self.ckpt_path)
self.model.load_state_dict(ckpt['model'])
else:
ckpt = torch.load(self.ckpt_path)
self.model.load_state_dict(ckpt['model_state_dict'])
self.start_epoch = 0
l = [d.name for d in os.scandir(self.val_path) if d.is_dir()]
l.sort()
l[l.index('Ram CV Cargo Van Minivan 2012'
)] = 'Ram C/V Cargo Van Minivan 2012'
self.label_texts = l
def main():
setup_default_logging()
args, args_text = _parse_args()
args.prefetcher = not args.no_prefetcher
args.distributed = False
if 'WORLD_SIZE' in os.environ:
args.distributed = int(os.environ['WORLD_SIZE']) > 1
if args.distributed and args.num_gpu > 1:
logging.warning(
'Using more than one GPU per process in distributed mode is not allowed. Setting num_gpu to 1.'
)
args.num_gpu = 1
args.device = 'cuda:0'
args.world_size = 1
args.rank = 0 # global rank
if args.distributed:
args.num_gpu = 1
args.device = 'cuda:%d' % args.local_rank
torch.cuda.set_device(args.local_rank)
torch.distributed.init_process_group(backend='nccl',
init_method='env://')
args.world_size = torch.distributed.get_world_size()
args.rank = torch.distributed.get_rank()
assert args.rank >= 0
if args.distributed:
logging.info(
'Training in distributed mode with multiple processes, 1 GPU per process. Process %d, total %d.'
% (args.rank, args.world_size))
else:
logging.info('Training with a single process on %d GPUs.' %
args.num_gpu)
torch.manual_seed(args.seed + args.rank)
model = create_model(args.model,
pretrained=args.pretrained,
num_classes=args.num_classes,
drop_rate=args.drop,
drop_connect_rate=args.drop_connect,
global_pool=args.gp,
bn_tf=args.bn_tf,
bn_momentum=args.bn_momentum,
bn_eps=args.bn_eps,
checkpoint_path=args.initial_checkpoint)
if args.local_rank == 0:
logging.info('Model %s created, param count: %d' %
(args.model, sum([m.numel()
for m in model.parameters()])))
data_config = resolve_data_config(vars(args),
model=model,
verbose=args.local_rank == 0)
num_aug_splits = 0
if args.aug_splits > 0:
assert args.aug_splits > 1, 'A split of 1 makes no sense'
num_aug_splits = args.aug_splits
if args.split_bn:
assert num_aug_splits > 1 or args.resplit
model = convert_splitbn_model(model, max(num_aug_splits, 2))
if args.num_gpu > 1:
if args.amp:
logging.warning(
'AMP does not work well with nn.DataParallel, disabling. Use distributed mode for multi-GPU AMP.'
)
args.amp = False
model = nn.DataParallel(model,
device_ids=list(range(args.num_gpu))).cuda()
else:
model.cuda()
optimizer = create_optimizer(args, model)
use_amp = False
if has_apex and args.amp:
model, optimizer = amp.initialize(model, optimizer, opt_level='O1')
use_amp = True
if args.local_rank == 0:
logging.info('NVIDIA APEX {}. AMP {}.'.format(
'installed' if has_apex else 'not installed',
'on' if use_amp else 'off'))
# optionally resume from a checkpoint
resume_state = {}
resume_epoch = None
if args.resume:
resume_state, resume_epoch = resume_checkpoint(model, args.resume)
if resume_state and not args.no_resume_opt:
if 'optimizer' in resume_state:
if args.local_rank == 0:
logging.info('Restoring Optimizer state from checkpoint')
optimizer.load_state_dict(resume_state['optimizer'])
if use_amp and 'amp' in resume_state and 'load_state_dict' in amp.__dict__:
if args.local_rank == 0:
logging.info('Restoring NVIDIA AMP state from checkpoint')
amp.load_state_dict(resume_state['amp'])
del resume_state
model_ema = None
if args.model_ema:
# Important to create EMA model after cuda(), DP wrapper, and AMP but before SyncBN and DDP wrapper
model_ema = ModelEma(model,
decay=args.model_ema_decay,
device='cpu' if args.model_ema_force_cpu else '',
resume=args.resume)
if args.distributed:
if args.sync_bn:
assert not args.split_bn
try:
if has_apex:
model = convert_syncbn_model(model)
else:
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(
model)
if args.local_rank == 0:
logging.info(
'Converted model to use Synchronized BatchNorm. WARNING: You may have issues if using '
'zero initialized BN layers (enabled by default for ResNets) while sync-bn enabled.'
)
except Exception as e:
logging.error(
'Failed to enable Synchronized BatchNorm. Install Apex or Torch >= 1.1'
)
if has_apex:
model = DDP(model, delay_allreduce=True)
else:
if args.local_rank == 0:
logging.info(
"Using torch DistributedDataParallel. Install NVIDIA Apex for Apex DDP."
)
model = DDP(model,
device_ids=[args.local_rank
]) # can use device str in Torch >= 1.1
# NOTE: EMA model does not need to be wrapped by DDP
lr_scheduler, num_epochs = create_scheduler(args, optimizer)
start_epoch = 0
if args.start_epoch is not None:
# a specified start_epoch will always override the resume epoch
start_epoch = args.start_epoch
elif resume_epoch is not None:
start_epoch = resume_epoch
if lr_scheduler is not None and start_epoch > 0:
lr_scheduler.step(start_epoch)
if args.local_rank == 0:
logging.info('Scheduled epochs: {}'.format(num_epochs))
train_dir = os.path.join(args.data, 'train')
if not os.path.exists(train_dir):
logging.error(
'Training folder does not exist at: {}'.format(train_dir))
exit(1)
dataset_train = Dataset(train_dir)
collate_fn = None
if args.prefetcher and args.mixup > 0:
assert not num_aug_splits # collate conflict (need to support deinterleaving in collate mixup)
collate_fn = FastCollateMixup(args.mixup, args.smoothing,
args.num_classes)
if num_aug_splits > 1:
dataset_train = AugMixDataset(dataset_train, num_splits=num_aug_splits)
loader_train = create_loader(
dataset_train,
input_size=data_config['input_size'],
batch_size=args.batch_size,
is_training=True,
use_prefetcher=args.prefetcher,
re_prob=args.reprob,
re_mode=args.remode,
re_count=args.recount,
re_split=args.resplit,
color_jitter=args.color_jitter,
auto_augment=args.aa,
num_aug_splits=num_aug_splits,
interpolation=args.train_interpolation,
mean=data_config['mean'],
std=data_config['std'],
num_workers=args.workers,
distributed=args.distributed,
collate_fn=collate_fn,
pin_memory=args.pin_mem,
)
eval_dir = os.path.join(args.data, 'val')
if not os.path.isdir(eval_dir):
eval_dir = os.path.join(args.data, 'validation')
if not os.path.isdir(eval_dir):
logging.error(
'Validation folder does not exist at: {}'.format(eval_dir))
exit(1)
dataset_eval = Dataset(eval_dir)
loader_eval = create_loader(
dataset_eval,
input_size=data_config['input_size'],
batch_size=4 * args.batch_size,
is_training=False,
use_prefetcher=args.prefetcher,
interpolation=data_config['interpolation'],
mean=data_config['mean'],
std=data_config['std'],
num_workers=args.workers,
distributed=args.distributed,
crop_pct=data_config['crop_pct'],
pin_memory=args.pin_mem,
)
if args.jsd:
assert num_aug_splits > 1 # JSD only valid with aug splits set
train_loss_fn = JsdCrossEntropy(num_splits=num_aug_splits,
smoothing=args.smoothing).cuda()
validate_loss_fn = nn.CrossEntropyLoss().cuda()
elif args.mixup > 0.:
# smoothing is handled with mixup label transform
train_loss_fn = SoftTargetCrossEntropy().cuda()
validate_loss_fn = nn.CrossEntropyLoss().cuda()
elif args.smoothing:
train_loss_fn = LabelSmoothingCrossEntropy(
smoothing=args.smoothing).cuda()
validate_loss_fn = nn.CrossEntropyLoss().cuda()
else:
train_loss_fn = nn.CrossEntropyLoss().cuda()
validate_loss_fn = train_loss_fn
eval_metric = args.eval_metric
best_metric = None
best_epoch = None
saver = None
output_dir = ''
if args.local_rank == 0:
output_base = args.output if args.output else './output'
exp_name = '-'.join([
datetime.now().strftime("%Y%m%d-%H%M%S"), args.model,
str(data_config['input_size'][-1])
])
output_dir = get_outdir(output_base, 'train', exp_name)
decreasing = True if eval_metric == 'loss' else False
saver = CheckpointSaver(checkpoint_dir=output_dir,
decreasing=decreasing)
with open(os.path.join(output_dir, 'args.yaml'), 'w') as f:
f.write(args_text)
try:
for epoch in range(start_epoch, num_epochs):
if args.distributed:
loader_train.sampler.set_epoch(epoch)
train_metrics = train_epoch(epoch,
model,
loader_train,
optimizer,
train_loss_fn,
args,
lr_scheduler=lr_scheduler,
saver=saver,
output_dir=output_dir,
use_amp=use_amp,
model_ema=model_ema)
if args.distributed and args.dist_bn in ('broadcast', 'reduce'):
if args.local_rank == 0:
logging.info(
"Distributing BatchNorm running means and vars")
distribute_bn(model, args.world_size, args.dist_bn == 'reduce')
eval_metrics = validate(model, loader_eval, validate_loss_fn, args)
if model_ema is not None and not args.model_ema_force_cpu:
if args.distributed and args.dist_bn in ('broadcast',
'reduce'):
distribute_bn(model_ema, args.world_size,
args.dist_bn == 'reduce')
ema_eval_metrics = validate(model_ema.ema,
loader_eval,
validate_loss_fn,
args,
log_suffix=' (EMA)')
eval_metrics = ema_eval_metrics
if lr_scheduler is not None:
# step LR for next epoch
lr_scheduler.step(epoch + 1, eval_metrics[eval_metric])
update_summary(epoch,
train_metrics,
eval_metrics,
os.path.join(output_dir, 'summary.csv'),
write_header=best_metric is None)
if saver is not None:
# save proper checkpoint with eval metric
save_metric = eval_metrics[eval_metric]
best_metric, best_epoch = saver.save_checkpoint(
model,
optimizer,
args,
epoch=epoch,
model_ema=model_ema,
metric=save_metric,
use_amp=use_amp)
except KeyboardInterrupt:
pass
if best_metric is not None:
logging.info('*** Best metric: {0} (epoch {1})'.format(
best_metric, best_epoch))
def validate(args):
_logger.info(f'\n\n ---------------EVALUATION {args.eps}------------------------------- \n\n')
_logger.info("Argument parser collected the following arguments:")
for arg in vars(args):
_logger.info(f" {arg}:{getattr(args, arg)}")
_logger.info("\n")
# might as well try to validate something
args.pretrained = args.pretrained or not args.checkpoint
args.prefetcher = not args.no_prefetcher
amp_autocast = suppress # do nothing
if args.amp:
if has_native_amp:
args.native_amp = True
elif has_apex:
args.apex_amp = True
else:
_logger.warning("Neither APEX or Native Torch AMP is available.")
assert not args.apex_amp or not args.native_amp, "Only one AMP mode should be set."
if args.native_amp:
amp_autocast = torch.cuda.amp.autocast
_logger.info('Validating in mixed precision with native PyTorch AMP.')
elif args.apex_amp:
_logger.info('Validating in mixed precision with NVIDIA APEX AMP.')
else:
_logger.info('Validating in float32. AMP not enabled.')
if args.legacy_jit:
set_jit_legacy()
# create model
model = create_model(
args.model,
pretrained=args.pretrained,
num_classes=args.num_classes,
in_chans=3,
global_pool=args.gp,
scriptable=args.torchscript)
if args.num_classes is None:
assert hasattr(model, 'num_classes'), 'Model must have `num_classes` attr if not set on cmd line/config.'
args.num_classes = model.num_classes
if args.checkpoint:
load_checkpoint(model, args.checkpoint, args.use_ema)
param_count = sum([m.numel() for m in model.parameters()])
_logger.info(
f'Model {args.model} created, param count: {param_count} ({(float(param_count)/(10.0**6)):.1f} M)'
)
data_config = resolve_data_config(vars(args), model=model, use_test_size=True, verbose=True)
test_time_pool = False
if not args.no_test_pool:
model, test_time_pool = apply_test_time_pool(model, data_config, use_test_size=True)
if args.torchscript:
torch.jit.optimized_execution(True)
model = torch.jit.script(model)
model = model.cuda()
if args.apex_amp:
model = amp.initialize(model, opt_level='O1')
if args.channels_last:
model = model.to(memory_format=torch.channels_last)
if args.num_gpu > 1:
model = torch.nn.DataParallel(model, device_ids=list(range(args.num_gpu)))
criterion = nn.CrossEntropyLoss().cuda()
dataset = create_dataset(
root=args.data_dir, name=args.dataset, split=args.split,
load_bytes=args.tf_preprocessing, class_map=args.class_map)
if args.valid_labels:
with open(args.valid_labels, 'r') as f:
valid_labels = {int(line.rstrip()) for line in f}
valid_labels = [i in valid_labels for i in range(args.num_classes)]
else:
valid_labels = None
if args.real_labels:
real_labels = RealLabelsImagenet(dataset.filenames(basename=True), real_json=args.real_labels)
else:
real_labels = None
crop_pct = 1.0 if test_time_pool else data_config['crop_pct']
loader = create_loader(
dataset,
input_size=data_config['input_size'],
batch_size=args.batch_size,
use_prefetcher=args.prefetcher,
interpolation=data_config['interpolation'],
mean=data_config['mean'],
std=data_config['std'],
num_workers=args.workers,
crop_pct=crop_pct,
pin_memory=args.pin_mem,
tf_preprocessing=args.tf_preprocessing)
batch_time = AverageMeter()
losses = AverageMeter()
top1 = AverageMeter()
top5 = AverageMeter()
top1_fgm_ae = AverageMeter()
top5_fgm_ae = AverageMeter()
top1_pgd_ae = AverageMeter()
top5_pgd_ae = AverageMeter()
model.eval()
#with torch.no_grad():# TODO Requires grad
# warmup, reduce variability of first batch time, especially for comparing torchscript vs non
input = torch.randn((args.batch_size,) + tuple(data_config['input_size'])).cuda()
if args.channels_last:
input = input.contiguous(memory_format=torch.channels_last)
model(input)
end = time.time()
for batch_idx, (input, target) in enumerate(loader):
if args.no_prefetcher:
target = target.cuda()
input = input.cuda()
if args.channels_last:
input = input.contiguous(memory_format=torch.channels_last)
# compute output
with amp_autocast():
output = model(input)
if valid_labels is not None:
output = output[:, valid_labels]
loss = criterion(output, target)
if real_labels is not None:
real_labels.add_result(output)
# TODO <---------------------
# Generate adversarial examples for current inputs
input_fgm_ae = fast_gradient_method(
model_fn=model,
x=input,
eps=args.eps,
norm=np.inf,
clip_min=None,
clip_max=None,
)
input_pgd_ae = projected_gradient_descent(
model_fn=model,
x=input,
eps=args.eps,
eps_iter=0.01,
nb_iter=40,
norm=np.inf,
clip_min=None,
clip_max=None,
)
# Predict with Adversarial Examples
with torch.no_grad():
with amp_autocast():
output_fgm_ae = model(input_fgm_ae)
output_pgd_ae = model(input_pgd_ae)
# measure accuracy and record loss
acc1, acc5 = accuracy(output.detach(), target, topk=(1, 5))
losses.update(loss.item(), input.size(0))
top1.update(acc1.item(), input.size(0))
top5.update(acc5.item(), input.size(0))
acc1_fgm_ae, acc5_fgm_ae = accuracy(output_fgm_ae.detach(), target, topk=(1, 5))
acc1_pgd_ae, acc5_pgd_ae = accuracy(output_pgd_ae.detach(), target, topk=(1, 5))
top1_fgm_ae.update(acc1_fgm_ae.item(), input.size(0))
top5_fgm_ae.update(acc5_fgm_ae.item(), input.size(0))
top1_pgd_ae.update(acc1_pgd_ae.item(), input.size(0))
top5_pgd_ae.update(acc5_pgd_ae.item(), input.size(0))
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if batch_idx % args.log_freq == 0:
_logger.info(
'Test: [{0:>4d}/{1}] '
'Time: {batch_time.val:.3f}s ({batch_time.avg:.3f}s, {rate_avg:>7.2f}/s) '
'Loss: {loss.val:>7.4f} ({loss.avg:>6.4f}) '
'Acc@1: {top1.val:>7.3f} ({top1.avg:>7.3f}) '
'Acc@5: {top5.val:>7.3f} ({top5.avg:>7.3f})'.format(
batch_idx, len(loader), batch_time=batch_time,
rate_avg=input.size(0) / batch_time.avg,
loss=losses, top1=top1, top5=top5))
if real_labels is not None:
raise NotImplementedError # TODO NOt modified for the adversarial examples mode
# real labels mode replaces topk values at the end
top1a, top5a = real_labels.get_accuracy(k=1), real_labels.get_accuracy(k=5)
else:
top1a, top5a = top1.avg, top5.avg
top1a_fgm_ae, top5a_fgm_ae = top1_fgm_ae.avg, top5_fgm_ae.avg
top1a_pgd_ae, top5a_pgd_ae = top1_pgd_ae.avg, top5_pgd_ae.avg
results = OrderedDict(
top1=round(top1a, 4), top1_err=round(100 - top1a, 4),
top5=round(top5a, 4), top5_err=round(100 - top5a, 4),
top1_fgm_ae=round(top1a_fgm_ae, 4),
top5_fgm_ae=round(top5a_fgm_ae, 4),
top1_pgd_ae=round(top1a_pgd_ae, 4),
top5_pgd_ae=round(top5a_pgd_ae, 4),
param_count=round(param_count / 1e6, 2),
img_size=data_config['input_size'][-1],
cropt_pct=crop_pct,
interpolation=data_config['interpolation'])
_logger.info(' * [Regular] Acc@1 {:.3f} ({:.3f}) Acc@5 {:.3f} ({:.3f})'.format(
results['top1'], results['top1_err'], results['top5'], results['top5_err']))
_logger.info(' * [FGM Adversarial Attack] Acc@1 {:.3f} Acc@5 {:.3f} '.format(
results['top1_fgm_ae'], results['top5_fgm_ae']))
_logger.info(' * [PGD Adversarial Attack] Acc@1 {:.3f} Acc@5 {:.3f} '.format(
results['top1_pgd_ae'], results['top5_pgd_ae']))
return results
