def which_network(network_name, **kwargs):
mean_std = None
# defining meand ans std
if kwargs['in_chns'] == 3:
mean, std = get_preprocessing_function('rgb')
mean_std = [mean, std]
elif kwargs['in_chns'] == 1:
mean, std = get_preprocessing_function('grey')
mean_std = [mean, std]
else:
mean, std = get_preprocessing_function(kwargs['in_chns'])
mean_std = [mean, std]
# checking whether it's an architecture or a network
if os.path.isfile(network_name):
checkpoint = torch.load(network_name, map_location='cpu')
architecture = checkpoint['arch']
if 'kwargs' in checkpoint:
kwargs = checkpoint['kwargs']
network = which_architecture(architecture, **kwargs)
network.load_state_dict(checkpoint['state_dict'])
if 'mean_std' in checkpoint:
mean_std = checkpoint['mean_std']
else:
network = which_architecture(network_name, **kwargs)
architecture = network_name
return network, architecture, mean_std
def get_transform(train,
colour_vision,
colour_space,
target_size=480,
other_tf=None):
# in the original version instead of target_size they use base_size of 520
min_size = int((0.5 if train else 1.0) * target_size)
max_size = int((2.0 if train else 1.0) * target_size)
transforms = [T.RandomResize(min_size, max_size)]
if train:
transforms.append(T.RandomHorizontalFlip(0.5))
transforms.append(T.RandomCrop(target_size))
if other_tf is not None:
transforms.append(*other_tf)
colour_transformation = preprocessing.colour_transformation(
colour_vision, colour_space)
# get ind 0 because it returns a list!
transforms.append(T.ColourTransformation(colour_transformation[0]))
transforms.append(T.ToTensor())
mean, std = get_preprocessing_function(colour_space, colour_vision)
transforms.append(T.Normalize(mean=mean, std=std))
return T.Compose(transforms)
def main(args):
args = parse_arguments(args)
if args.model_path is None:
args.model_path = args.model_name
colour_space = args.colour_space
target_size = args.target_size
# loading the model
is_segmentation = False
if 'deeplabv3_resnet' in args.model_name or 'fcn_resnet' in args.model_name:
is_segmentation = True
transfer_weights = [args.model_path, None, is_segmentation]
model = pretrained_models.get_pretrained_model(args.model_name,
transfer_weights)
# selecting the layer
model = pretrained_models.LayerActivation(
pretrained_models.get_backbones(args.model_name, model),
args.activation_layer)
model = model.eval()
model.cuda()
mean, std = model_utils.get_preprocessing_function(colour_space,
'trichromat')
transform = torch_transforms.Compose([
cv2_transforms.ToTensor(),
cv2_transforms.Normalize(mean, std),
])
if args.task == '2afc':
default_dist = DISTORTIONS_2AFC
db_class = image_quality.BAPPS2afc
run_fun = run_2afc
else:
default_dist = DISTORTIONS_JND
db_class = image_quality.BAPPSjnd
run_fun = run_jnd
distortions = default_dist if args.distortion is None else [
args.distortion
]
eval_results = dict()
for dist in distortions:
print('Starting with %s' % dist)
db = db_class(root=args.db_dir,
split=args.split,
distortion=dist,
transform=transform)
db_loader = torch.utils.data.DataLoader(db,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
print_val = dist if args.print else None
eval_results[dist] = run_fun(db_loader, model, print_val)
save_results(eval_results, args.out_file)
def main(args):
args = parse_arguments(args)
if args.model_path is None:
args.model_path = args.model_name
colour_space = args.colour_space
target_size = args.target_size
# loading the model
is_segmentation = False
if 'deeplabv3_resnet' in args.model_name or 'fcn_resnet' in args.model_name:
is_segmentation = True
transfer_weights = [args.model_path, None, is_segmentation]
model = pretrained_models.get_pretrained_model(args.model_name,
transfer_weights)
# selecting the layer
model = pretrained_models.LayerActivation(
pretrained_models.get_backbones(args.model_name, model),
args.activation_layer)
model = model.eval()
model.cuda()
mean, std = model_utils.get_preprocessing_function(colour_space,
'trichromat')
transform = torch_transforms.Compose([
cv2_transforms.ToTensor(),
cv2_transforms.Normalize(mean, std),
])
parts = PARTS if args.part is None else [args.part]
eval_results = dict()
for part in parts:
print('Starting with %s' % part)
db = image_quality.LIVE(root=args.db_dir,
part=part,
transform=transform)
db_loader = torch.utils.data.DataLoader(db,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
print_val = part if args.print else None
eval_results[part] = run_live(db_loader, model, print_val)
save_results(eval_results, args.out_file)
def main_worker(ngpus_per_node, args):
mean, std = model_utils.get_preprocessing_function(args.colour_space,
args.vision_type)
if args.gpus is not None:
print("Use GPU: {} for training".format(args.gpus))
if args.distributed:
if args.dist_url == "env://" and args.rank == -1:
args.rank = int(os.environ["RANK"])
if args.multiprocessing_distributed:
# For multiprocessing distributed training, rank needs to be the
# global rank among all the processes
args.rank = args.rank * ngpus_per_node + args.gpus
dist.init_process_group(backend=args.dist_backend,
init_method=args.dist_url,
world_size=args.world_size,
rank=args.rank)
# create model
if args.transfer_weights is not None:
print('Transferred model!')
(model, _) = model_utils.which_network(args.transfer_weights[0],
args.task_type,
num_classes=args.old_classes)
which_layer = -1
if len(args.transfer_weights) == 2:
which_layer = args.transfer_weights[1]
model = model_utils.NewClassificationModel(model, which_layer,
args.num_classes)
elif args.custom_arch:
print('Custom model!')
supported_customs = ['resnet_basic_custom', 'resnet_bottleneck_custom']
if os.path.isfile(args.network_name):
checkpoint = torch.load(args.network_name, map_location='cpu')
customs = None
if 'customs' in checkpoint:
customs = checkpoint['customs']
# TODO: num_classes is just for backward compatibility
if 'num_classes' not in customs:
customs['num_classes'] = 1000
model = which_architecture(checkpoint['arch'], customs,
args.contrast_head)
args.network_name = checkpoint['arch']
model.load_state_dict(checkpoint['state_dict'], strict=False)
elif args.network_name in supported_customs:
model = custom_models.__dict__[args.network_name](
args.blocks,
contrast_head=args.contrast_head,
pooling_type=args.pooling_type,
in_chns=len(mean),
num_classes=args.num_classes,
inplanes=args.num_kernels,
kernel_size=args.kernel_size)
elif args.pretrained:
print("=> using pre-trained model '{}'".format(args.network_name))
model = models.__dict__[args.network_name](pretrained=True)
else:
print("=> creating model '{}'".format(args.network_name))
model = models.__dict__[args.network_name]()
# TODO: why load weights is False?
args.out_dir = prepare_training.prepare_output_directories(
dataset_name='contrast',
network_name=args.network_name,
optimiser='sgd',
load_weights=False,
experiment_name=args.experiment_name,
framework='pytorch')
# preparing the output folder
create_dir(args.out_dir)
json_file_name = os.path.join(args.out_dir, 'args.json')
with open(json_file_name, 'w') as fp:
json.dump(dict(args._get_kwargs()), fp, sort_keys=True, indent=4)
if args.distributed:
# For multiprocessing distributed, DistributedDataParallel constructor
# should always set the single device scope, otherwise,
# DistributedDataParallel will use all available devices.
if args.gpus is not None:
torch.cuda.set_device(args.gpus)
model.cuda(args.gpus)
# When using a single GPU per process and per
# DistributedDataParallel, we need to divide the batch size
# ourselves based on the total number of GPUs we have
args.batch_size = int(args.batch_size / ngpus_per_node)
args.workers = int(args.workers / ngpus_per_node)
model = torch.nn.parallel.DistributedDataParallel(
model, device_ids=[args.gpus])
else:
model.cuda()
# DistributedDataParallel will divide and allocate batch_size to all
# available GPUs if device_ids are not set
model = torch.nn.parallel.DistributedDataParallel(model)
elif args.gpus is not None:
torch.cuda.set_device(args.gpus)
model = model.cuda(args.gpus)
else:
# DataParallel will divide and allocate batch_size to all available GPUs
if (args.network_name.startswith('alexnet')
or args.network_name.startswith('vgg')):
model.features = torch.nn.DataParallel(model.features)
model.cuda()
else:
model = torch.nn.DataParallel(model).cuda()
# define loss function (criterion) and optimizer
criterion = nn.CrossEntropyLoss().cuda(args.gpus)
# optimiser
if args.transfer_weights is None:
optimizer = torch.optim.SGD(model.parameters(),
args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
else:
# for p in model.features.parameters():
# p.requires_grad = False
params_to_optimize = [
{
'params': [p for p in model.features.parameters()],
'lr': 1e-6
},
{
'params': [p for p in model.fc.parameters()]
},
]
optimizer = torch.optim.SGD(params_to_optimize,
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
model_progress = []
model_progress_path = os.path.join(args.out_dir, 'model_progress.csv')
# optionally resume from a checkpoint
# TODO: it would be best if resume load the architecture from this file
# TODO: merge with which_architecture
best_acc1 = 0
if args.resume is not None:
if os.path.isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume, map_location='cpu')
args.initial_epoch = checkpoint['epoch']
best_acc1 = checkpoint['best_acc1']
model.load_state_dict(checkpoint['state_dict'])
if args.gpus is not None:
# best_acc1 may be from a checkpoint from a different GPU
best_acc1 = best_acc1.to(args.gpus)
model = model.cuda(args.gpus)
optimizer.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
if os.path.exists(model_progress_path):
model_progress = np.loadtxt(model_progress_path, delimiter=',')
model_progress = model_progress.tolist()
else:
print("=> no checkpoint found at '{}'".format(args.resume))
cudnn.benchmark = True
train_trans = []
valid_trans = []
both_trans = []
if args.mosaic_pattern is not None:
mosaic_trans = preprocessing.mosaic_transformation(args.mosaic_pattern)
both_trans.append(mosaic_trans)
if args.num_augmentations != 0:
augmentations = preprocessing.random_augmentation(
args.augmentation_settings, args.num_augmentations)
train_trans.append(augmentations)
target_size = default_configs.get_default_target_size(
args.dataset, args.target_size)
# loading the training set
train_trans = [*both_trans, *train_trans]
db_params = {
'colour_space': args.colour_space,
'vision_type': args.vision_type,
'mask_image': args.mask_image
}
if args.dataset in ['imagenet', 'celeba', 'natural']:
path_or_sample = args.data_dir
else:
path_or_sample = args.train_samples
train_dataset = dataloader.train_set(args.dataset,
target_size,
mean,
std,
extra_transformation=train_trans,
data_dir=path_or_sample,
**db_params)
if args.dataset == 'natural':
train_dataset.num_crops = args.batch_size
args.batch_size = 1
if args.distributed:
train_sampler = torch.utils.data.distributed.DistributedSampler(
train_dataset)
else:
train_sampler = None
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=(train_sampler is None),
num_workers=args.workers,
pin_memory=True,
sampler=train_sampler)
# loading validation set
valid_trans = [*both_trans, *valid_trans]
validation_dataset = dataloader.validation_set(
args.dataset,
target_size,
mean,
std,
extra_transformation=valid_trans,
data_dir=path_or_sample,
**db_params)
if args.dataset == 'natural':
validation_dataset.num_crops = train_dataset.num_crops
args.batch_size = 1
val_loader = torch.utils.data.DataLoader(validation_dataset,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
# training on epoch
for epoch in range(args.initial_epoch, args.epochs):
if args.distributed:
train_sampler.set_epoch(epoch)
misc_utils.adjust_learning_rate(optimizer, epoch, args)
# train for one epoch
train_log = train_on_data(train_loader, model, criterion, optimizer,
epoch, args)
# evaluate on validation set
validation_log = validate_on_data(val_loader, model, criterion, args)
model_progress.append([*train_log, *validation_log])
# remember best acc@1 and save checkpoint
acc1 = validation_log[2]
is_best = acc1 > best_acc1
best_acc1 = max(acc1, best_acc1)
if misc_utils.is_saving_node(args.multiprocessing_distributed,
args.rank, ngpus_per_node):
misc_utils.save_checkpoint(
{
'epoch': epoch + 1,
'arch': args.network_name,
'customs': {
'pooling_type': args.pooling_type,
'in_chns': len(mean),
'num_classes': args.num_classes,
'blocks': args.blocks,
'num_kernels': args.num_kernels,
'kernel_size': args.kernel_size
},
'preprocessing': {
'mean': mean,
'std': std
},
'state_dict': model.state_dict(),
'best_acc1': best_acc1,
'optimizer': optimizer.state_dict(),
'target_size': target_size,
},
is_best,
out_folder=args.out_dir)
# TODO: get this header directly as a dictionary keys
header = 'epoch,t_time,t_loss,t_top1,t_top5,v_time,v_loss,v_top1,v_top5'
np.savetxt(model_progress_path,
np.array(model_progress),
delimiter=',',
header=header)
def main(args):
args = parse_arguments(args)
if args.imagenet_dir is None:
args.imagenet_dir = '/home/arash/Software/imagenet/raw-data/validation/'
colour_space = args.colour_space
target_size = args.target_size
if args.repeat:
target_size = int(args.target_size / 2)
mean, std = model_utils.get_preprocessing_function(colour_space,
'trichromat')
extra_transformations = []
if args.noise is not None:
noise_kwargs = {'amount': float(args.noise[1])}
extra_transformations.append(
cv2_preprocessing.UniqueTransformation(imutils.gaussian_noise,
**noise_kwargs))
if args.mosaic_pattern is not None:
mosaic_trans = cv2_preprocessing.MosaicTransformation(
args.mosaic_pattern)
extra_transformations.append(mosaic_trans)
# testing setting
freqs = args.freqs
if freqs is None:
if args.model_fest:
test_sfs = [
107.558006181068, 60.2277887428434, 42.5666842683747,
30.1176548488805, 21.2841900445655, 15.0589946730047,
10.6611094334144, 7.5293662203794, 5.33055586478039,
4.01568430264343
]
args.gabor = 'model_fest'
else:
if target_size == 256:
t4s = [
target_size / 2,
target_size / 2.5,
target_size / 3,
target_size / 3.5,
target_size / 3.75,
target_size / 4,
]
else:
# assuming 128
t4s = [64]
sf_base = ((target_size / 2) / np.pi)
test_sfs = [
sf_base / e for e in [
*np.arange(1, 21), *np.arange(21, 61, 5),
*np.arange(61, t4s[-1], 25), *t4s
]
]
else:
if len(freqs) == 3:
test_sfs = np.linspace(freqs[0], freqs[1], int(freqs[2]))
else:
test_sfs = freqs
# so the sfs gets sorted
test_sfs = np.unique(test_sfs)
contrasts = args.contrasts
if contrasts is None:
test_contrasts = [0.5]
else:
test_contrasts = contrasts
test_thetas = np.linspace(0, np.pi, 7)
test_rhos = np.linspace(0, np.pi, 4)
test_ps = [0.0, 1.0]
if args.pretrained:
if args.side_by_side:
if args.scale_factor is None:
scale_factor = (args.target_size / 256)**2
else:
scale_factor = args.scale_factor
model = pretrained_models.NewClassificationModel(
args.model_path,
grey_width=args.grey_width == 40,
scale_factor=scale_factor)
else:
if args.scale_factor is None:
scale_factor = (args.target_size / 128)**2
else:
scale_factor = args.scale_factor
model = models_csf.ContrastDiscrimination(
args.model_path,
grey_width=args.grey_width == 40,
scale_factor=scale_factor)
else:
model, _ = model_utils.which_network_classification(args.model_path, 2)
model.eval()
model.cuda()
mean_std = None
if args.visualise:
mean_std = (mean, std)
if args.avg_illuminant < 0:
max_high = 1 + 2 * args.avg_illuminant
elif args.avg_illuminant > 0:
max_high = 1 + -2 * args.avg_illuminant
else:
max_high = 1.0
mid_contrast = (0 + max_high) / 2
all_results = None
csf_flags = [mid_contrast for _ in test_sfs]
if args.db == 'gratings':
for i in range(len(csf_flags)):
low = 0
high = max_high
j = 0
while csf_flags[i] is not None:
print('%.2d %.3d Doing %f - %f %f %f' %
(i, j, test_sfs[i], csf_flags[i], low, high))
test_samples = {
'amp': [csf_flags[i]],
'lambda_wave': [test_sfs[i]],
'theta': test_thetas,
'rho': test_rhos,
'side': test_ps,
'avg_illuminant': args.avg_illuminant
}
db_params = {
'colour_space': colour_space,
'vision_type': args.vision_type,
'repeat': args.repeat,
'mask_image': args.gabor,
'grey_width': args.grey_width,
'side_by_side': args.side_by_side
}
db = dataloader.validation_set(args.db,
target_size,
mean,
std,
extra_transformations,
data_dir=test_samples,
**db_params)
db.contrast_space = args.contrast_space
db_loader = torch.utils.data.DataLoader(
db,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True,
)
if args.side_by_side:
new_results, all_results = run_gratings(
db_loader,
model,
args.out_file,
args.print,
mean_std=mean_std,
old_results=all_results)
else:
new_results, all_results = run_gratings_separate(
db_loader,
model,
args.out_file,
args.print,
mean_std=mean_std,
old_results=all_results)
new_contrast, low, high = sensitivity_sf(new_results,
test_sfs[i],
varname='all',
th=0.75,
low=low,
high=high)
if (abs(csf_flags[i] - max_high) < 1e-3
or new_contrast == csf_flags[i] or j == 20):
print('had to skip', csf_flags[i])
csf_flags[i] = None
else:
csf_flags[i] = new_contrast
j += 1
def main_worker(ngpus_per_node, args):
mean, std = model_utils.get_preprocessing_function(args.colour_space,
args.vision_type)
# preparing the output folder
create_dir(args.out_dir)
if args.gpus is not None:
print("Use GPU: {} for training".format(args.gpus))
if args.distributed:
if args.dist_url == "env://" and args.rank == -1:
args.rank = int(os.environ["RANK"])
if args.multiprocessing_distributed:
# For multiprocessing distributed training, rank needs to be the
# global rank among all the processes
args.rank = args.rank * ngpus_per_node + args.gpus
dist.init_process_group(backend=args.dist_backend,
init_method=args.dist_url,
world_size=args.world_size,
rank=args.rank)
# create model
if args.transfer_weights is not None:
print('Transferred model!')
model = contrast_utils.AFCModel(args.network_name,
args.transfer_weights)
elif args.custom_arch:
print('Custom model!')
supported_customs = ['resnet_basic_custom', 'resnet_bottleneck_custom']
if args.network_name in supported_customs:
model = custom_models.__dict__[args.network_name](
args.blocks,
pooling_type=args.pooling_type,
in_chns=len(mean),
num_classes=args.num_classes,
inplanes=args.num_kernels,
kernel_size=args.kernel_size)
elif args.pretrained:
print("=> using pre-trained model '{}'".format(args.network_name))
model = models.__dict__[args.network_name](pretrained=True)
else:
print("=> creating model '{}'".format(args.network_name))
model = models.__dict__[args.network_name]()
if args.distributed:
# For multiprocessing distributed, DistributedDataParallel constructor
# should always set the single device scope, otherwise,
# DistributedDataParallel will use all available devices.
if args.gpus is not None:
torch.cuda.set_device(args.gpus)
model.cuda(args.gpus)
# When using a single GPU per process and per
# DistributedDataParallel, we need to divide the batch size
# ourselves based on the total number of GPUs we have
args.batch_size = int(args.batch_size / ngpus_per_node)
args.workers = int(args.workers / ngpus_per_node)
model = torch.nn.parallel.DistributedDataParallel(
model, device_ids=[args.gpus])
else:
model.cuda()
# DistributedDataParallel will divide and allocate batch_size to all
# available GPUs if device_ids are not set
model = torch.nn.parallel.DistributedDataParallel(model)
elif args.gpus is not None:
torch.cuda.set_device(args.gpus)
model = model.cuda(args.gpus)
else:
# DataParallel will divide and allocate batch_size to all available GPUs
if (args.network_name.startswith('alexnet')
or args.network_name.startswith('vgg')):
model.features = torch.nn.DataParallel(model.features)
model.cuda()
else:
model = torch.nn.DataParallel(model).cuda()
# define loss function (criterion) and optimizer
criterion = soft_cross_entropy
# optimiser
if args.transfer_weights is None:
optimizer = torch.optim.SGD(model.parameters(),
args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
else:
params_to_optimize = [
{
'params': [p for p in model.parameters() if p.requires_grad]
},
]
optimizer = torch.optim.SGD(params_to_optimize,
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
model_progress = []
model_progress_path = os.path.join(args.out_dir, 'model_progress.csv')
# optionally resume from a checkpoint
# TODO: it would be best if resume load the architecture from this file
# TODO: merge with which_architecture
best_acc1 = 0
if args.resume is not None:
if os.path.isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume, map_location='cpu')
args.initial_epoch = checkpoint['epoch']
best_acc1 = checkpoint['best_acc1']
model.load_state_dict(checkpoint['state_dict'])
if args.gpus is not None:
# best_acc1 may be from a checkpoint from a different GPU
best_acc1 = best_acc1.to(args.gpus)
model = model.cuda(args.gpus)
optimizer.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
if os.path.exists(model_progress_path):
model_progress = np.loadtxt(model_progress_path, delimiter=',')
model_progress = model_progress.tolist()
else:
print("=> no checkpoint found at '{}'".format(args.resume))
cudnn.benchmark = True
train_trans = []
valid_trans = []
both_trans = []
if args.mosaic_pattern is not None:
mosaic_trans = preprocessing.mosaic_transformation(args.mosaic_pattern)
both_trans.append(mosaic_trans)
if args.num_augmentations != 0:
augmentations = preprocessing.random_augmentation(
args.augmentation_settings, args.num_augmentations)
train_trans.append(augmentations)
target_size = default_configs.get_default_target_size(
args.dataset, args.target_size)
final_trans = [
cv2_transforms.ToTensor(),
cv2_transforms.Normalize(mean, std),
]
train_trans.append(
cv2_transforms.RandomResizedCrop(target_size, scale=(0.08, 1.0)))
# loading the training set
train_trans = torch_transforms.Compose(
[*both_trans, *train_trans, *final_trans])
train_dataset = image_quality.BAPPS2afc(root=args.data_dir,
split='train',
transform=train_trans,
concat=0.5)
if args.distributed:
train_sampler = torch.utils.data.distributed.DistributedSampler(
train_dataset)
else:
train_sampler = None
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=(train_sampler is None),
num_workers=args.workers,
pin_memory=True,
sampler=train_sampler)
valid_trans.extend([
cv2_transforms.Resize(target_size),
cv2_transforms.CenterCrop(target_size),
])
# loading validation set
valid_trans = torch_transforms.Compose(
[*both_trans, *valid_trans, *final_trans])
validation_dataset = image_quality.BAPPS2afc(root=args.data_dir,
split='val',
transform=valid_trans,
concat=0)
val_loader = torch.utils.data.DataLoader(validation_dataset,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
# training on epoch
for epoch in range(args.initial_epoch, args.epochs):
if args.distributed:
train_sampler.set_epoch(epoch)
misc_utils.adjust_learning_rate(optimizer, epoch, args)
# train for one epoch
train_log = train_on_data(train_loader, model, criterion, optimizer,
epoch, args)
# evaluate on validation set
validation_log = validate_on_data(val_loader, model, criterion, args)
model_progress.append([*train_log, *validation_log])
# remember best acc@1 and save checkpoint
acc1 = validation_log[2]
is_best = acc1 > best_acc1
best_acc1 = max(acc1, best_acc1)
if misc_utils.is_saving_node(args.multiprocessing_distributed,
args.rank, ngpus_per_node):
misc_utils.save_checkpoint(
{
'epoch': epoch + 1,
'arch': args.network_name,
'customs': {
'pooling_type': args.pooling_type,
'in_chns': len(mean),
'num_classes': args.num_classes,
'blocks': args.blocks,
'num_kernels': args.num_kernels,
'kernel_size': args.kernel_size
},
'transfer_weights': args.transfer_weights,
'preprocessing': {
'mean': mean,
'std': std
},
'state_dict': model.state_dict(),
'best_acc1': best_acc1,
'optimizer': optimizer.state_dict(),
'target_size': target_size,
},
is_best,
out_folder=args.out_dir)
# TODO: get this header directly as a dictionary keys
header = 'epoch,t_time,t_loss,t_top5,v_time,v_loss,v_top1'
np.savetxt(model_progress_path,
np.array(model_progress),
delimiter=',',
header=header)
def main(args):
args.gpus = set_visible_gpus(args.gpus)
args.device = prepare_device(args.gpus)
args.architecture_kwargs = {
'in_chns': args.in_chns
}
# TODO: this is not a nice solution here
if 'wavenet' in args.architecture:
args.architecture_kwargs['inplanes'] = args.num_kernels
args.architecture_kwargs['planes'] = args.blocks
# creating the model
model, architecture, mean_std = geetup_net.which_network(
args.architecture, **args.architecture_kwargs
)
model = model.to(args.device)
args.out_dir = prepare_training.prepare_output_directories(
dataset_name='geetup_' + args.dataset, network_name=architecture,
optimiser='sgd', load_weights=False,
experiment_name=args.experiment_name, framework='pytorch'
)
logging.basicConfig(
filename=args.out_dir + '/experiment_info.log', filemode='w',
format='%(levelname)s: %(message)s', level=logging.INFO
)
validation_pickle = _get_multi_sensory_paths(
args.data_dir, args.validation_file
)
validation_dataset = geetup_db.get_validation_dataset(
validation_pickle, args.target_size, mean_std
)
validation_loader = torch.utils.data.DataLoader(
validation_dataset, batch_size=args.batch_size, shuffle=False,
num_workers=args.workers, pin_memory=True
)
if args.random is not None:
mean, std = get_preprocessing_function('rgb', 'trichromat')
normalize_inverse = NormalizeInverse(mean, std)
process_random_image(model, validation_loader, normalize_inverse, args)
return
args.criterion = nn.BCELoss().to(args.device)
if args.evaluate:
predict_outs = predict(
validation_loader, model, args.criterion, args
)
for key, item in predict_outs.items():
file_name = '_'.join(e for e in [key, *args.out_prefix])
result_file = '%s/%s_%s' % (
args.out_dir, file_name, args.validation_file
)
write_pickle(result_file, item)
return
training_pickle = _get_multi_sensory_paths(
args.data_dir, args.train_file
)
train_dataset = geetup_db.get_train_dataset(
training_pickle, args.target_size, mean_std, args.crop_scale,
args.gaussian_sigma
)
train_loader = torch.utils.data.DataLoader(
train_dataset, batch_size=args.batch_size, shuffle=True,
num_workers=args.workers, pin_memory=True
)
# optimiser
optimizer = torch.optim.SGD(
model.parameters(), args.lr,
momentum=args.momentum, weight_decay=args.weight_decay
)
epochs(model, train_loader, validation_loader, optimizer, args)
def main_worker(ngpus_per_node, args):
global best_acc1
is_pill_img = 'wcs_xyz_png_1600' in args.data_dir
mean, std = model_utils.get_preprocessing_function(args.colour_space,
args.vision_type)
# preparing the output folder
create_dir(args.out_dir)
if args.gpus is not None:
print("Use GPU: {} for training".format(args.gpus))
if args.distributed:
if args.dist_url == "env://" and args.rank == -1:
args.rank = int(os.environ["RANK"])
if args.multiprocessing_distributed:
# For multiprocessing distributed training, rank needs to be the
# global rank among all the processes
args.rank = args.rank * ngpus_per_node + args.gpus
dist.init_process_group(backend=args.dist_backend,
init_method=args.dist_url,
world_size=args.world_size,
rank=args.rank)
# create model
if args.prediction:
checkpoint = torch.load(args.network_weights, map_location='cpu')
blocks = checkpoint['customs']['blocks']
pooling_type = checkpoint['customs']['pooling_type']
num_kernels = checkpoint['customs']['num_kernels']
outputs = checkpoint['customs']['outputs']
for key, val in outputs.items():
if 'area' not in val:
outputs[key] = None
model = resnet.__dict__[args.network_name](blocks,
pooling_type=pooling_type,
in_chns=len(mean),
inplanes=num_kernels,
outputs=outputs)
model.load_state_dict(checkpoint['state_dict'])
elif args.transfer_weights is not None:
print('Transferred model!')
(model, _) = model_utils.which_network(args.transfer_weights,
args.task_type,
num_classes=args.old_classes)
model = model_utils.NewClassificationModel(model, args.num_classes)
elif args.custom_arch:
print('Custom model!')
if (args.network_name == 'resnet_basic_custom'
or args.network_name == 'resnet_bottleneck_custom'):
outputs = {'objects': None, 'munsells': None, 'illuminants': None}
imagenet_weights = args.imagenet_weights
if args.object_area is not None:
outputs['objects'] = {
'num_classes': 2100,
'area': args.object_area
}
if args.munsell_area is not None:
outputs['munsells'] = {
'num_classes': 1600,
'area': args.munsell_area
}
if args.illuminant_area is not None:
outputs['illuminants'] = {
'num_classes': 280,
'area': args.illuminant_area
}
model = resnet.__dict__[args.network_name](
args.blocks,
pooling_type=args.pooling_type,
in_chns=len(mean),
inplanes=args.num_kernels,
outputs=outputs,
imagenet_weights=imagenet_weights)
elif args.pretrained:
print("=> using pre-trained model '{}'".format(args.network_name))
model = models.__dict__[args.network_name](pretrained=True)
else:
print("=> creating model '{}'".format(args.network_name))
model = models.__dict__[args.network_name]()
if args.distributed:
# For multiprocessing distributed, DistributedDataParallel constructor
# should always set the single device scope, otherwise,
# DistributedDataParallel will use all available devices.
if args.gpus is not None:
torch.cuda.set_device(args.gpus)
model.cuda(args.gpus)
# When using a single GPU per process and per
# DistributedDataParallel, we need to divide the batch size
# ourselves based on the total number of GPUs we have
args.batch_size = int(args.batch_size / ngpus_per_node)
args.workers = int(args.workers / ngpus_per_node)
model = torch.nn.parallel.DistributedDataParallel(
model, device_ids=[args.gpus])
else:
model.cuda()
# DistributedDataParallel will divide and allocate batch_size to all
# available GPUs if device_ids are not set
model = torch.nn.parallel.DistributedDataParallel(model)
elif args.gpus is not None:
torch.cuda.set_device(args.gpus)
model = model.cuda(args.gpus)
else:
# DataParallel will divide and allocate batch_size to all available GPUs
if (args.network_name.startswith('alexnet')
or args.network_name.startswith('vgg')):
model.features = torch.nn.DataParallel(model.features)
model.cuda()
else:
model = torch.nn.DataParallel(model).cuda()
# define loss function (criterion) and optimizer
criterion = nn.CrossEntropyLoss().cuda(args.gpus)
# optimiser
optimizer = torch.optim.SGD(model.parameters(),
args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
model_progress = []
model_progress_path = os.path.join(args.out_dir, 'model_progress.csv')
# optionally resume from a checkpoint
# TODO: it would be best if resume load the architecture from this file
# TODO: merge with which_architecture
if args.resume is not None:
if os.path.isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume, map_location='cpu')
args.initial_epoch = checkpoint['epoch']
best_acc1 = checkpoint['best_acc1']
model.load_state_dict(checkpoint['state_dict'])
if args.gpus is not None:
# best_acc1 may be from a checkpoint from a different GPU
best_acc1 = best_acc1.to(args.gpus)
model = model.cuda(args.gpus)
optimizer.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
if os.path.exists(model_progress_path):
model_progress = np.loadtxt(model_progress_path, delimiter=',')
model_progress = model_progress.tolist()
else:
print("=> no checkpoint found at '{}'".format(args.resume))
cudnn.benchmark = True
normalize = transforms.Normalize(mean=mean, std=std)
other_transformations = []
if args.num_augmentations != 0:
augmentations = preprocessing.random_augmentation(
args.augmentation_settings, args.num_augmentations)
other_transformations.append(augmentations)
target_size = get_default_target_size(args.dataset, args.target_size)
train_dataset, validation_dataset = get_train_val_dataset(
args.data_dir, other_transformations, [], normalize,
args.imagenet_weights)
if args.distributed:
train_sampler = torch.utils.data.distributed.DistributedSampler(
train_dataset)
else:
train_sampler = None
if args.prediction:
manipulation_values = args.parameters['kwargs'][args.manipulation]
manipulation_name = args.parameters['f_name']
for j, manipulation_value in enumerate(manipulation_values):
args.parameters['kwargs'][args.manipulation] = manipulation_value
prediction_transformation = preprocessing.prediction_transformation(
args.parameters, args.colour_space,
tmp_c_space(manipulation_name))
other_transformations = [prediction_transformation]
_, validation_dataset = get_train_val_dataset(
args.data_dir, other_transformations, other_transformations,
normalize, args.imagenet_weights)
val_loader = torch.utils.data.DataLoader(
validation_dataset,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
pred_log = predict(val_loader, model, criterion,
torch.device(args.gpus))
from kernelphysiology.dl.utils import prepapre_testing
prepapre_testing.save_predictions(pred_log, args.experiment_name,
args.pred_name, args.dataset,
manipulation_name,
manipulation_value)
return
val_loader = torch.utils.data.DataLoader(validation_dataset,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=(train_sampler is None),
num_workers=args.workers,
pin_memory=True,
sampler=train_sampler)
if args.ill_colour is not None:
print('Performing with illuminant correction')
args.ill_colour = np.loadtxt(args.ill_colour, delimiter=',')
# training on epoch
for epoch in range(args.initial_epoch, args.epochs):
if args.distributed:
train_sampler.set_epoch(epoch)
if args.imagenet_weights is None:
adjust_learning_rate(optimizer, epoch, args)
# train for one epoch
train_log = train_on_data(train_loader, model, criterion, optimizer,
epoch, args)
# evaluate on validation set
validation_log = validate_on_data(val_loader, model, criterion, args)
model_progress.append([*train_log, *validation_log])
# remember best acc@1 and save checkpoint
acc1 = validation_log[2]
is_best = acc1 > best_acc1
best_acc1 = max(acc1, best_acc1)
if not args.multiprocessing_distributed or (
args.multiprocessing_distributed
and args.rank % ngpus_per_node == 0):
save_checkpoint(
{
'epoch': epoch + 1,
'arch': args.network_name,
'customs': {
'pooling_type': args.pooling_type,
'in_chns': len(mean),
'num_classes': args.num_classes,
'blocks': args.blocks,
'num_kernels': args.num_kernels,
'outputs': outputs
},
'state_dict': model.state_dict(),
'best_acc1': best_acc1,
'optimizer': optimizer.state_dict(),
'target_size': target_size,
},
is_best,
out_folder=args.out_dir)
# TODO: get this header directly as a dictionary keys
header = 'epoch,t_time,t_loss,t_lo,t_lm,t_li,t_ao,t_am,t_ai,' \
'v_time,v_loss,v_lo,v_lm,v_li,v_ao,v_am,v_ai'
np.savetxt(model_progress_path,
np.array(model_progress),
delimiter=',',
header=header)
def train_on_data(train_loader, model, criterion, optimizer, epoch, args):
losses = AverageMeter()
batch_time = AverageMeter()
data_time = AverageMeter()
losses_obj = AverageMeter()
top1_obj = AverageMeter()
top5_obj = AverageMeter()
losses_mun = AverageMeter()
top1_mun = AverageMeter()
top5_mun = AverageMeter()
losses_ill = AverageMeter()
top1_ill = AverageMeter()
top5_ill = AverageMeter()
if args.top_k is None:
topks = (1, )
else:
topks = (1, args.top_k)
# switch to train mode
model.train()
mean, std = model_utils.get_preprocessing_function(args.colour_space,
args.vision_type)
normalise_inverse = cv2_transforms.NormalizeInverse(mean, std)
normalise_back = transforms.Normalize(mean=mean, std=std)
end = time.time()
for i, (input_image, targets) in enumerate(train_loader):
# measure data loading time
data_time.update(time.time() - end)
if args.gpus is not None:
input_image = input_image.cuda(args.gpus, non_blocking=True)
targets = targets.cuda(args.gpus, non_blocking=True)
# compute output
out_obj, out_mun, out_ill = model(input_image)
if out_obj is None:
loss_obj = 0
else:
loss_obj = criterion(out_obj, targets[:, 0])
acc1_obj, acc5_obj = accuracy(out_obj, targets[:, 0], topk=topks)
losses_obj.update(loss_obj.item(), input_image.size(0))
top1_obj.update(acc1_obj[0], input_image.size(0))
top5_obj.update(acc5_obj[0], input_image.size(0))
if out_mun is None:
loss_mun = 0
else:
loss_mun = criterion(out_mun, targets[:, 1])
acc1_mun, acc5_mun = accuracy(out_mun, targets[:, 1], topk=topks)
losses_mun.update(loss_mun.item(), input_image.size(0))
top1_mun.update(acc1_mun[0], input_image.size(0))
top5_mun.update(acc5_mun[0], input_image.size(0))
if out_ill is None:
loss_ill = 0
else:
loss_ill = criterion(out_ill, targets[:, 2])
acc1_ill, acc5_ill = accuracy(out_ill, targets[:, 2], topk=topks)
losses_ill.update(loss_ill.item(), input_image.size(0))
top1_ill.update(acc1_ill[0], input_image.size(0))
top5_ill.update(acc5_ill[0], input_image.size(0))
loss = loss_obj + loss_mun + loss_ill
losses.update(loss.item(), input_image.size(0))
# compute gradient and do SGD step
optimizer.zero_grad()
loss.backward()
optimizer.step()
if out_mun is None and args.ill_colour is not None:
input_image2 = correct_image(normalise_inverse, normalise_back,
input_image, out_ill, args.ill_colour)
out_obj2, out_mun2, _ = model(input_image2)
loss_mun2 = 0
loss_obj2 = 0
if out_mun2 is not None:
loss_mun2 = criterion(out_mun2, targets[:, 1])
if out_obj2 is not None:
loss_obj2 = criterion(out_obj2, targets[:, 0])
loss2 = loss_obj2 + loss_mun2
optimizer.zero_grad()
loss2.backward()
optimizer.step()
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
# printing the accuracy at certain intervals
if i % args.print_freq == 0:
print('Epoch: [{0}][{1}/{2}]\t'
'Time {batch_time.val:.2f} ({batch_time.avg:.2f})\t'
'Data {data_time.val:.2f} ({data_time.avg:.2f})\t'
'Loss {loss.val:.2f} ({loss.avg:.2f})\t'
'LO {obj_loss.val:.2f} ({obj_loss.avg:.2f})\t'
'LM {mun_loss.val:.2f} ({mun_loss.avg:.2f})\t'
'LI {ill_loss.val:.2f} ({ill_loss.avg:.2f})\t'
'Ao {obj_acc.val:.2f} ({obj_acc.avg:.2f})\t'
'AM {mun_acc.val:.2f} ({mun_acc.avg:.2f})\t'
'AI {ill_acc.val:.2f} ({ill_acc.avg:.2f})'.format(
epoch,
i,
len(train_loader),
batch_time=batch_time,
data_time=data_time,
loss=losses,
obj_loss=losses_obj,
mun_loss=losses_mun,
ill_loss=losses_ill,
obj_acc=top1_obj,
mun_acc=top1_mun,
ill_acc=top1_ill))
return [
epoch, batch_time.avg, losses.avg, losses_obj.avg, losses_mun.avg,
losses_ill.avg, top1_obj.avg, top1_mun.avg, top1_ill.avg
]
def generic_evaluation(args, fn, save_fn=None, **kwargs):
manipulation_values = args.parameters['kwargs'][args.manipulation]
manipulation_name = args.parameters['f_name']
other_mans = args.parameters['others']
for j, current_network in enumerate(args.network_files):
# which architecture
(model, target_size) = model_utils.which_network(
current_network,
args.task_type,
num_classes=args.num_classes,
kill_kernels=args.kill_kernels,
kill_planes=args.kill_planes,
kill_lines=args.kill_lines)
model.to(args.device)
mean, std = model_utils.get_preprocessing_function(
args.colour_space, args.network_chromaticities[j])
normalize = transforms.Normalize(mean=mean, std=std)
for i, manipulation_value in enumerate(manipulation_values):
args.parameters['kwargs'][args.manipulation] = manipulation_value
output_file = prepapre_testing._prepare_saving_file(
args.experiment_name,
args.network_names[j],
args.dataset,
manipulation_name,
manipulation_value,
extension='csv')
if os.path.exists(output_file):
continue
if args.task_type == 'segmentation' or 'voc' in args.dataset:
prediction_transformation = preprocessing.prediction_transformation_seg(
args.parameters, args.colour_space,
tmp_c_space(manipulation_name))
else:
prediction_transformation = preprocessing.prediction_transformation(
args.parameters, args.colour_space,
tmp_c_space(manipulation_name))
colour_vision = 'trichromat'
if _requires_colour_transform(manipulation_name,
args.network_chromaticities[j]):
colour_vision = args.network_chromaticities[j]
other_transformations = []
for oth_man in other_mans:
if args.task_type == 'segmentation' or 'voc' in args.dataset:
other_transformations.append(
preprocessing.prediction_transformation_seg(
oth_man, args.colour_space,
tmp_c_space(oth_man['f_name'])))
else:
other_transformations.append(
preprocessing.prediction_transformation(
oth_man, args.colour_space,
tmp_c_space(oth_man['f_name'])))
if args.mosaic_pattern is not None:
other_transformations.append(
preprocessing.mosaic_transformation(args.mosaic_pattern))
if args.sf_filter is not None:
other_transformations.append(
preprocessing.sf_transformation(args.sf_filter,
args.sf_filter_chn))
other_transformations.append(prediction_transformation)
print('Processing network %s and %s %f' %
(current_network, manipulation_name, manipulation_value))
# which dataset
# reading it after the model, because each might have their own
# specific size
# loading validation set
target_size = get_default_target_size(args.dataset,
args.target_size)
target_transform = utils_db.ImagenetCategoryTransform(
args.categories, args.cat_dir)
validation_dataset = utils_db.get_validation_dataset(
args.dataset,
args.validation_dir,
colour_vision,
args.colour_space,
other_transformations,
normalize,
target_size,
task=args.task_type,
target_transform=target_transform)
# TODO: nicer solution:
if 'sampler' not in args:
sampler = None
else:
sampler = args.sampler(validation_dataset)
if 'collate_fn' not in args:
args.collate_fn = None
# FIXME: add segmentation datasests
val_loader = torch.utils.data.DataLoader(
validation_dataset,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True,
sampler=sampler,
collate_fn=args.collate_fn)
if args.random_images is not None:
out_folder = prepapre_testing.prepare_saving_dir(
args.experiment_name, args.network_names[j], args.dataset,
manipulation_name)
normalize_inverse = NormalizeInverse(mean, std)
fn(val_loader, out_folder, normalize_inverse,
manipulation_value, **kwargs)
elif args.activation_map is not None:
model = model_utils.LayerActivation(model, args.activation_map)
current_results = fn(val_loader, model, **kwargs)
save_fn(current_results, args.experiment_name,
args.network_names[j], args.dataset, manipulation_name,
manipulation_value)
else:
(_, _, current_results) = fn(val_loader, model, **kwargs)
save_fn(current_results, args.experiment_name,
args.network_names[j], args.dataset, manipulation_name,
manipulation_value)