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 get_val_dataset(val_dir, target_size, preprocess):
mean, std = preprocess
normalise = cv2_transforms.Normalize(mean=mean, std=std)
transform = torch_transforms.Compose([
cv2_transforms.Resize(target_size),
cv2_transforms.CenterCrop(target_size),
cv2_transforms.ToTensor(),
normalise,
])
train_dataset = ImageFolder({'root': val_dir, 'transform': transform})
return train_dataset
def get_train_dataset(train_dir, target_size, preprocess):
mean, std = preprocess
normalise = cv2_transforms.Normalize(mean=mean, std=std)
scale = (0.08, 1.0)
size_transform = cv2_transforms.RandomResizedCrop(target_size, scale=scale)
transform = torch_transforms.Compose([
size_transform,
cv2_transforms.RandomHorizontalFlip(),
cv2_transforms.ToTensor(),
normalise,
])
train_dataset = ImageFolder({'root': train_dir, 'transform': transform})
return train_dataset
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(args):
args = parse_arguments(args)
if args.random_seed < 0:
os.environ['PYTHONHASHSEED'] = str(args.random_seed)
torch.manual_seed(args.random_seed)
torch.cuda.manual_seed_all(args.random_seed)
torch.cuda.manual_seed(args.random_seed)
np.random.seed(args.random_seed)
random.seed(args.random_seed)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
weights_rgb = torch.load(args.model_path, map_location='cpu')
network = vqmodel.VQ_CVAE(128,
k=args.k,
kl=args.kl,
in_chns=3,
cos_distance=args.cos_dis)
network.load_state_dict(weights_rgb)
if args.exclude > 0:
which_vec = [args.exclude - 1]
print(which_vec)
network.state_dict()['emb.weight'][:, which_vec] = 0
elif args.exclude < 0:
which_vec = [*range(8)]
which_vec.remove(abs(args.exclude) - 1)
print(which_vec)
network.state_dict()['emb.weight'][:, which_vec] = 0
network.cuda()
network.eval()
if not os.path.exists(args.out_dir):
os.mkdir(args.out_dir)
args.in_colour_space = args.colour_space[:3]
args.out_colour_space = args.colour_space[4:7]
(imagenet_model, target_size) = model_utils.which_network(
args.imagenet_model,
'classification',
num_classes=args.num_classes,
)
imagenet_model.cuda()
imagenet_model.eval()
mean = (0.5, 0.5, 0.5)
std = (0.5, 0.5, 0.5)
if 'vggface2' in args.validation_dir:
transform_funcs = transforms.Compose([
cv2_transforms.Resize(256),
cv2_transforms.CenterCrop(224),
cv2_transforms.ToTensor(),
cv2_transforms.Normalize(mean, std)
])
imagenet_transformations = transforms.Compose([
cv2_transforms.ToTensor(),
cv2_transforms.Normalize([0.485, 0.456, 0.406],
[0.229, 0.224, 0.225])
])
else:
transform_funcs = transforms.Compose([
cv2_transforms.Resize(512),
cv2_transforms.CenterCrop(512),
cv2_transforms.ToTensor(),
cv2_transforms.Normalize(mean, std)
])
imagenet_transformations = transforms.Compose([
cv2_transforms.Resize(256),
cv2_transforms.CenterCrop(224),
cv2_transforms.ToTensor(),
cv2_transforms.Normalize([0.485, 0.456, 0.406],
[0.229, 0.224, 0.225])
])
intransform_funs = []
if args.in_colour_space != ' rgb':
intransform_funs.append(
cv2_preprocessing.VisionTypeTransformation(None,
args.in_colour_space))
if args.noise[0] is not None:
value = float(args.noise[1])
noise_name = args.noise[0]
if noise_name == 'sp':
noise_fun = imutils.s_p_noise
kwargs = {'amount': value, 'seed': args.random_seed}
elif noise_name == 'gaussian':
noise_fun = imutils.gaussian_noise
kwargs = {'amount': value, 'seed': args.random_seed}
elif noise_name == 'speckle':
noise_fun = imutils.speckle_noise
kwargs = {'amount': value, 'seed': args.random_seed}
elif noise_name == 'blur':
noise_fun = imutils.gaussian_blur
kwargs = {'sigmax': value, 'seed': args.random_seed}
if noise_name != 'blur':
kwargs['eq_chns'] = True
intransform_funs.append(
cv2_preprocessing.UniqueTransformation(noise_fun, **kwargs))
intransform = transforms.Compose(intransform_funs)
test_loader = torch.utils.data.DataLoader(ImageFolder(
root=args.validation_dir,
intransform=intransform,
outtransform=None,
transform=transform_funcs),
batch_size=args.batch_size,
shuffle=False)
top1, top5, prediction_output = export(test_loader, network, mean, std,
imagenet_model,
imagenet_transformations, args)
output_file = '%s/%s.csv' % (args.out_dir, args.colour_space)
np.savetxt(output_file, prediction_output, delimiter=',', fmt='%i')
def _get_shared_post_transforms(mean, std):
return [
cv2_transforms.ToTensor(),
cv2_transforms.Normalize(mean, std),
]
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, _) = 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 args.network_name in supported_customs:
model = custom_models.__dict__[args.network_name](
args.blocks,
pooling_type=args.pooling_type,
in_chns=len(mean * args.extra_chns),
num_classes=args.num_classes,
inplanes=args.num_kernels)
else:
model = custom_models.__dict__[args.network_name](
pooling_type=args.pooling_type,
in_chns=len(mean * args.extra_chns),
num_classes=args.num_classes)
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
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.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
if args.extra_chns > 1:
normalize_train = RandomNormalize(mean=mean,
std=std,
extra_chns=args.extra_chns)
else:
normalize_train = cv2_transforms.Normalize(mean=mean, std=std)
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)
# loading the training set
train_trans = [*both_trans, *train_trans]
train_dataset = utils_db.get_train_dataset(args.dataset, args.train_dir,
args.vision_type,
args.colour_space, train_trans,
normalize_train, target_size)
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)
if args.extra_chns > 1:
normalize_val = RandomNormalize(mean=mean,
std=std,
extra_chns=args.extra_chns,
is_val=True)
else:
normalize_val = cv2_transforms.Normalize(mean=mean, std=std)
# loading validation set
valid_trans = [*both_trans, *valid_trans]
validation_dataset = utils_db.get_validation_dataset(
args.dataset,
args.validation_dir,
args.vision_type,
args.colour_space,
valid_trans,
normalize_val,
target_size,
)
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 = misc_utils.train_on_data(train_loader, model, criterion,
optimizer, epoch, args)
# evaluate on validation set
validation_log = misc_utils.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 * args.extra_chns),
'num_classes': args.num_classes,
'blocks': args.blocks,
'num_kernels': args.num_kernels
},
'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)
weights_rgb = torch.load(args.model_path, map_location='cpu')
network = vqmodel.VQ_CVAE(128, k=args.k, kl=args.kl, in_chns=3,
cos_distance=args.cos_dis)
network.load_state_dict(weights_rgb)
if args.exclude > 0:
which_vec = [args.exclude - 1]
print(which_vec)
network.state_dict()['emb.weight'][:, which_vec] = 0
elif args.exclude < 0:
which_vec = [*range(8)]
which_vec.remove(abs(args.exclude) - 1)
print(which_vec)
network.state_dict()['emb.weight'][:, which_vec] = 0
network.cuda()
network.eval()
if not os.path.exists(args.out_dir):
os.mkdir(args.out_dir)
intransform_funs = []
mean = (0.5, 0.5, 0.5)
std = (0.5, 0.5, 0.5)
manipulation_func = []
args.suffix = ''
if args.manipulation is not None:
man_func = None
parameters = None
if args.manipulation[0] == 'contrast':
man_func = imutils.adjust_contrast
parameters = {'amount': float(args.manipulation[1])}
elif args.manipulation[0] == 'gamma':
man_func = imutils.adjust_gamma
parameters = {'amount': float(args.manipulation[1])}
elif args.manipulation[0] == 'luminance':
man_func = imutils.reduce_lightness
parameters = {'amount': float(args.manipulation[1])}
elif args.manipulation[0] == 'chromaticity':
man_func = imutils.reduce_chromaticity
parameters = {'amount': float(args.manipulation[1])}
elif args.manipulation[0] == 'red_green':
man_func = imutils.reduce_red_green
parameters = {'amount': float(args.manipulation[1])}
elif args.manipulation[0] == 'yellow_blue':
man_func = imutils.reduce_yellow_blue
parameters = {'amount': float(args.manipulation[1])}
elif args.manipulation[0] == 'illuminant':
man_func = imutils.adjust_illuminant
parameters = {'illuminant': [
float(args.manipulation[1]),
float(args.manipulation[2]),
float(args.manipulation[3])
]}
if man_func is None:
sys.exit('Unsupported function %s' % args.manipulation[0])
args.suffix = '_' + args.manipulation[0] + '_' + ''.join(
e for e in args.manipulation[1:]
)
manipulation_func.append(cv2_preprocessing.UniqueTransformation(
man_func, **parameters
))
intransform_funs.append(*manipulation_func)
args.in_colour_space = args.colour_space[:3]
args.out_colour_space = args.colour_space[4:7]
if args.in_colour_space != ' rgb':
intransform_funs.append(
cv2_preprocessing.ColourSpaceTransformation(args.in_colour_space)
)
intransform = transforms.Compose(intransform_funs)
transform_funcs = transforms.Compose([
# cv2_transforms.Resize(256), cv2_transforms.CenterCrop(224),
cv2_transforms.ToTensor(),
cv2_transforms.Normalize(mean, std)
])
if args.dataset == 'imagenet':
test_loader = torch.utils.data.DataLoader(
data_loaders.ImageFolder(
root=args.validation_dir,
intransform=intransform,
outtransform=None,
transform=transform_funcs
),
batch_size=args.batch_size, shuffle=False
)
elif args.dataset == 'celeba':
test_loader = torch.utils.data.DataLoader(
data_loaders.CelebA(
root=args.validation_dir,
intransform=intransform,
outtransform=None,
transform=transform_funcs,
split='test'
),
batch_size=args.batch_size, shuffle=False
)
else:
test_loader = torch.utils.data.DataLoader(
data_loaders.CategoryImages(
root=args.validation_dir,
# FIXME
category=args.category,
intransform=intransform,
outtransform=None,
transform=transform_funcs
),
batch_size=args.batch_size, shuffle=False
)
export(test_loader, network, mean, std, args)
def main(args):
args = parse_arguments(args)
args.cuda = not args.no_cuda and torch.cuda.is_available()
if args.random_seed is not None:
os.environ['PYTHONHASHSEED'] = str(args.random_seed)
torch.manual_seed(args.random_seed)
torch.cuda.manual_seed_all(args.random_seed)
torch.cuda.manual_seed(args.random_seed)
np.random.seed(args.random_seed)
random.seed(args.random_seed)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
if args.cuda:
args.gpus = [int(i) for i in args.gpus.split(',')]
torch.cuda.set_device(args.gpus[0])
cudnn.benchmark = True
args.mean = (0.5, 0.5, 0.5)
args.std = (0.5, 0.5, 0.5)
if 'labhue' in args.colour_space:
args.mean = (0.5, 0.5, 0.5, 0.5)
args.std = (0.5, 0.5, 0.5, 0.5)
normalise = transforms.Normalize(args.mean, args.std)
lr = args.lr or default_hyperparams[args.dataset]['lr']
k = args.k or default_hyperparams[args.dataset]['k']
hidden = args.hidden or default_hyperparams[args.dataset]['hidden']
num_channels = args.num_channels or dataset_n_channels[args.dataset]
target_size = args.target_size or dataset_target_size[args.dataset]
dataset_transforms = {
'custom': transforms.Compose(
[transforms.Resize(256), transforms.CenterCrop(224),
transforms.ToTensor(), normalise]),
'coco': transforms.Compose([normalise]),
'voc': transforms.Compose(
[cv2_transforms.RandomResizedCropSegmentation(target_size,
scale=(0.50, 1.0)),
cv2_transforms.ToTensorSegmentation(),
cv2_transforms.NormalizeSegmentation(args.mean, args.std)]),
'bsds': transforms.Compose(
[cv2_transforms.RandomResizedCropSegmentation(target_size,
scale=(0.50, 1.0)),
cv2_transforms.ToTensorSegmentation(),
cv2_transforms.NormalizeSegmentation(args.mean, args.std)]),
'imagenet': transforms.Compose(
[cv2_transforms.Resize(target_size + 32),
cv2_transforms.CenterCrop(target_size),
cv2_transforms.ToTensor(),
cv2_transforms.Normalize(args.mean, args.std)]),
'celeba': transforms.Compose(
[cv2_transforms.Resize(target_size + 32),
cv2_transforms.CenterCrop(target_size),
cv2_transforms.ToTensor(),
cv2_transforms.Normalize(args.mean, args.std)]),
'cifar10': transforms.Compose(
[transforms.ToTensor(), normalise]),
'mnist': transforms.ToTensor()
}
save_path = vae_util.setup_logging_from_args(args)
writer = SummaryWriter(save_path)
in_colour_space = args.colour_space[:3]
out_colour_space = args.colour_space[4:]
args.colour_space = out_colour_space
if args.model == 'wavenet':
# model = wavenet_vae.wavenet_bottleneck(
# latent_dim=k, in_channels=num_channels
# )
task = None
out_chns = 3
if 'voc' in args.dataset:
task = 'segmentation'
out_chns = 21
from torchvision.models import resnet
backbone = resnet.__dict__['resnet50'](
pretrained=True,
replace_stride_with_dilation=[False, True, True]
)
from torchvision.models._utils import IntermediateLayerGetter
return_layers = {'layer4': 'out'}
resnet = IntermediateLayerGetter(
backbone, return_layers=return_layers
)
model = vae_model.ResNet_VQ_CVAE(
hidden, k=k, resnet=resnet, num_channels=num_channels,
colour_space=args.colour_space, task=task,
out_chns=out_chns
)
elif args.model == 'vae':
model = vanilla_vae.VanillaVAE(latent_dim=args.k, in_channels=3)
else:
task = None
out_chns = 3
if 'voc' in args.dataset:
task = 'segmentation'
out_chns = 21
elif 'bsds' in args.dataset:
task = 'segmentation'
out_chns = 1
elif args.colour_space == 'labhue':
out_chns = 4
backbone = None
if args.backbone is not None:
backbone = {
'arch_name': args.backbone[0],
'layer_name': args.backbone[1]
}
if len(args.backbone) > 2:
backbone['weights_path'] = args.backbone[2]
models[args.dataset][args.model] = vae_model.Backbone_VQ_VAE
model = models[args.dataset][args.model](
hidden, k=k, kl=args.kl, num_channels=num_channels,
colour_space=args.colour_space, task=task,
out_chns=out_chns, cos_distance=args.cos_dis,
use_decor_loss=args.decor, backbone=backbone
)
if args.cuda:
model.cuda()
if args.load_encoder is not None:
params_to_optimize = [
{'params': [p for p in model.decoder.parameters() if
p.requires_grad]},
{'params': [p for p in model.fc.parameters() if
p.requires_grad]},
]
optimizer = optim.Adam(params_to_optimize, lr=lr)
else:
optimizer = optim.Adam(model.parameters(), lr=lr)
scheduler = optim.lr_scheduler.StepLR(
optimizer, int(args.epochs / 3), 0.5
)
if args.resume is not None:
checkpoint = torch.load(args.resume, map_location='cpu')
model.load_state_dict(checkpoint['state_dict'])
model.cuda()
args.start_epoch = checkpoint['epoch'] + 1
scheduler.load_state_dict(checkpoint['scheduler'])
optimizer.load_state_dict(checkpoint['optimizer'])
elif args.fine_tune is not None:
weights = torch.load(args.fine_tune, map_location='cpu')
model.load_state_dict(weights, strict=False)
model.cuda()
elif args.load_encoder is not None:
weights = torch.load(args.load_encoder, map_location='cpu')
weights = removekey(weights, 'decode')
model.load_state_dict(weights, strict=False)
model.cuda()
intransform_funs = []
if args.gamma is not None:
kwargs = {'amount': args.gamma}
augmentation_settings = [
{'function': imutils.adjust_gamma, 'kwargs': kwargs}
]
intransform_funs.append(
cv2_preprocessing.RandomAugmentationTransformation(
augmentation_settings, num_augmentations=1
)
)
if args.mosaic_pattern is not None:
intransform_funs.append(
cv2_preprocessing.MosaicTransformation(args.mosaic_pattern)
)
if in_colour_space != 'rgb':
intransform_funs.append(
cv2_preprocessing.ColourSpaceTransformation(in_colour_space)
)
intransform = transforms.Compose(intransform_funs)
outtransform_funs = []
args.inv_func = None
if args.colour_space is not None:
outtransform_funs.append(
cv2_preprocessing.ColourSpaceTransformation(args.colour_space)
)
if args.vis_rgb:
args.inv_func = lambda x: generic_inv_fun(x, args.colour_space)
outtransform = transforms.Compose(outtransform_funs)
if args.data_dir is not None:
args.train_dir = os.path.join(args.data_dir, 'train')
args.validation_dir = os.path.join(args.data_dir, 'validation')
else:
args.train_dir = args.train_dir
args.validation_dir = args.validation_dir
kwargs = {'num_workers': args.workers,
'pin_memory': True} if args.cuda else {}
args.vis_func = vae_util.grid_save_reconstructed_images
if args.colour_space == 'labhue':
args.vis_func = vae_util.grid_save_reconstructed_labhue
if args.dataset == 'coco':
train_loader = panoptic_utils.get_coco_train(
args.batch_size, args.opts, args.cfg_file
)
test_loader = panoptic_utils.get_coco_test(
args.batch_size, args.opts, args.cfg_file
)
elif 'voc' in args.dataset:
train_loader = torch.utils.data.DataLoader(
data_loaders.VOCSegmentation(
root=args.data_dir,
image_set='train',
intransform=intransform,
outtransform=outtransform,
transform=dataset_transforms[args.dataset],
**dataset_train_args[args.dataset]
),
batch_size=args.batch_size, shuffle=True, **kwargs
)
test_loader = torch.utils.data.DataLoader(
data_loaders.VOCSegmentation(
root=args.data_dir,
image_set='val',
intransform=intransform,
outtransform=outtransform,
transform=dataset_transforms[args.dataset],
**dataset_test_args[args.dataset]
),
batch_size=args.batch_size, shuffle=False, **kwargs
)
elif args.category is not None:
train_loader = torch.utils.data.DataLoader(
data_loaders.CategoryImages(
root=args.train_dir,
category=args.category,
intransform=intransform,
outtransform=outtransform,
transform=dataset_transforms[args.dataset],
**dataset_train_args[args.dataset]
),
batch_size=args.batch_size, shuffle=True, **kwargs
)
test_loader = torch.utils.data.DataLoader(
data_loaders.CategoryImages(
root=args.validation_dir,
category=args.category,
intransform=intransform,
outtransform=outtransform,
transform=dataset_transforms[args.dataset],
**dataset_test_args[args.dataset]
),
batch_size=args.batch_size, shuffle=False, **kwargs
)
elif args.dataset == 'bsds':
args.vis_func = vae_util.grid_save_reconstructed_bsds
train_loader = torch.utils.data.DataLoader(
datasets_classes[args.dataset](
root=args.data_dir,
intransform=intransform,
outtransform=outtransform,
transform=dataset_transforms[args.dataset],
**dataset_train_args[args.dataset]
),
batch_size=args.batch_size, shuffle=True, **kwargs
)
test_loader = torch.utils.data.DataLoader(
datasets_classes[args.dataset](
root=args.data_dir,
intransform=intransform,
outtransform=outtransform,
transform=dataset_transforms[args.dataset],
**dataset_test_args[args.dataset]
),
batch_size=args.batch_size, shuffle=False, **kwargs
)
elif args.dataset == 'celeba':
train_loader = torch.utils.data.DataLoader(
datasets_classes[args.dataset](
root=args.data_dir,
intransform=intransform,
outtransform=outtransform,
transform=dataset_transforms[args.dataset],
split='train',
**dataset_train_args[args.dataset]
),
batch_size=args.batch_size, shuffle=True, **kwargs
)
test_loader = torch.utils.data.DataLoader(
datasets_classes[args.dataset](
root=args.data_dir,
intransform=intransform,
outtransform=outtransform,
transform=dataset_transforms[args.dataset],
split='test',
**dataset_test_args[args.dataset]
),
batch_size=args.batch_size, shuffle=False, **kwargs
)
else:
train_loader = torch.utils.data.DataLoader(
datasets_classes[args.dataset](
root=args.train_dir,
intransform=intransform,
outtransform=outtransform,
transform=dataset_transforms[args.dataset],
**dataset_train_args[args.dataset]
),
batch_size=args.batch_size, shuffle=True, **kwargs
)
test_loader = torch.utils.data.DataLoader(
datasets_classes[args.dataset](
root=args.validation_dir,
intransform=intransform,
outtransform=outtransform,
transform=dataset_transforms[args.dataset],
**dataset_test_args[args.dataset]
),
batch_size=args.batch_size, shuffle=False, **kwargs
)
console = logging.StreamHandler()
console.setLevel(logging.INFO)
logging.getLogger('').addHandler(console)
for epoch in range(args.start_epoch, args.epochs):
train_losses = train(
epoch, model, train_loader, optimizer, args.cuda, args.log_interval,
save_path, args
)
test_losses = test_net(epoch, model, test_loader, args.cuda, save_path,
args)
for k in train_losses.keys():
name = k.replace('_train', '')
train_name = k
test_name = k.replace('train', 'test')
writer.add_scalars(
name, {'train': train_losses[train_name],
'test': test_losses[test_name]}, epoch
)
scheduler.step()
vae_util.save_checkpoint(
{
'epoch': epoch,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict(),
'scheduler': scheduler.state_dict(),
'arch': {'k': args.k, 'hidden': args.hidden}
},
save_path
)
def main(args):
args = parse_arguments(args)
weights_net = torch.load(args.model_path, map_location='cpu')
args.in_colour_space = args.colour_space[:3]
args.out_colour_space = args.colour_space[4:]
args.outs_dict = dict()
args.outs_dict[args.out_colour_space] = {'shape': [1, 1, 3]}
from segmentation_models import unet
network = unet.model.Unet(in_channels=3,
encoder_weights=None,
outs_dict=args.outs_dict,
classes=3)
network.load_state_dict(weights_net)
if args.exclude > 0:
which_vec = [args.exclude - 1]
print(which_vec)
network.state_dict()['emb.weight'][:, which_vec] = 0
elif args.exclude < 0:
which_vec = [*range(8)]
which_vec.remove(abs(args.exclude) - 1)
print(which_vec)
network.state_dict()['emb.weight'][:, which_vec] = 0
network.cuda()
network.eval()
if not os.path.exists(args.out_dir):
os.mkdir(args.out_dir)
mean = (0.5, 0.5, 0.5)
std = (0.5, 0.5, 0.5)
transform_funcs = transforms.Compose([
cv2_transforms.Resize(args.target_size + 32),
cv2_transforms.CenterCrop(args.target_size),
cv2_transforms.ToTensor(),
cv2_transforms.Normalize(mean, std)
])
intransform_funs = []
if args.in_colour_space != ' rgb':
intransform_funs.append(
cv2_preprocessing.ColourSpaceTransformation(args.in_colour_space))
if args.corr_noise:
parameters = dict()
parameters['function'] = imutils.s_p_noise
parameters['kwargs'] = {'amount': args.corr_noise, 'eq_chns': True}
intransform_funs.append(
cv2_preprocessing.PredictionTransformation(parameters))
intransform = transforms.Compose(intransform_funs)
if args.dataset == 'imagenet':
test_loader = torch.utils.data.DataLoader(data_loaders.ImageFolder(
root=args.validation_dir,
intransform=intransform,
outtransform=None,
transform=transform_funcs),
batch_size=args.batch_size,
shuffle=False)
elif args.dataset == 'celeba':
test_loader = torch.utils.data.DataLoader(data_loaders.CelebA(
root=args.validation_dir,
intransform=intransform,
outtransform=None,
transform=transform_funcs,
split='test'),
batch_size=args.batch_size,
shuffle=False)
else:
test_loader = torch.utils.data.DataLoader(
data_loaders.CategoryImages(
root=args.validation_dir,
# FIXME
category=args.category,
intransform=intransform,
outtransform=None,
transform=transform_funcs),
batch_size=args.batch_size,
shuffle=False)
export(test_loader, network, mean, std, args)
def main(argv):
parser = argparse.ArgumentParser("")
parser.add_argument('-aname',
'--architecture',
required=True,
type=str,
help='Name of the architecture or network')
parser.add_argument('--data_dir',
type=str,
default=None,
help='The path to the data directory (default: None)')
args = parser.parse_args(argv)
# creating the model
model = pmodels.__dict__[args.architecture](pretrained=True).cuda()
model.eval()
# creating the hooks
act_dict, rfhs = _create_resnet_hooks(model)
# transformations
mean = [0.485, 0.456, 0.406]
std = [0.229, 0.224, 0.225]
transform = torch_transforms.Compose([
cv2_transforms.ToTensor(),
cv2_transforms.Normalize(mean, std),
])
# preparing the data
data_file_path = '%s/simulationOrder_preliminary.csv' % args.data_dir
data_file = np.loadtxt(data_file_path, delimiter=',', dtype=str)
data_file = data_file[1:]
test_colour = data_file[:, 4].astype('int')
shape = data_file[:, 5].astype('int')
n_comps = 4
img_dir = '%s/img' % (args.data_dir)
all_preds = []
# looping over all the data
for test_ind, current_test in enumerate(data_file):
print('Doing [%d/%d]' % (test_ind, len(data_file)))
img_target = io.imread('%s/%s' % (img_dir, current_test[0]))
img_ref0 = io.imread('%s/%s' % (img_dir, current_test[1]))
img_ref1 = io.imread('%s/%s' % (img_dir, current_test[2]))
img_ref2 = io.imread('%s/%s' % (img_dir, current_test[3]))
# making it pytorch friendly
img_batch = torch.stack([
transform(img_target),
transform(img_ref0),
transform(img_ref1),
transform(img_ref2)
]).cuda()
# computing the activations
net_out = model(img_batch)
act_dict['fc'] = net_out.detach()
with torch.no_grad():
layer_pred = []
for key, val in act_dict.items():
current_acts = val.clone().cpu().numpy().squeeze()
key_distance_mat = np.zeros((n_comps, n_comps))
for i in range(0, n_comps - 1):
for j in range(i, n_comps):
if i == j:
continue
diff = current_acts[i] - current_acts[j]
diff = (diff**2).sum()**0.5
key_distance_mat[i, j] = diff
key_distance_mat[j, i] = diff
distance_mat = np.mean(key_distance_mat, axis=0)
layer_pred.append(distance_mat.argmax())
all_preds.append(layer_pred)
# saving the results
out_file = '%s.csv' % args.architecture
np.savetxt(out_file, np.array(all_preds), delimiter=',')
def main(args):
args = arguments.parse_arguments(args)
# determining the number of input channels
args.in_chns = 3
out_chns = 3
args.out_chns = out_chns
args.mean = 0.5
args.std = 0.5
target_size = args.target_size or dataset_target_size[args.dataset]
if args.dataset == 'ccvr':
pre_shared_transforms = [
cv2_transforms.Resize(target_size + 32),
cv2_transforms.RandomCrop(target_size),
]
else:
pre_shared_transforms = [
cv2_transforms.Resize(target_size + 32),
cv2_transforms.CenterCrop(target_size),
]
post_shared_transforms = [
cv2_transforms.ToTensor(),
cv2_transforms.Normalize(args.mean, args.std)
]
pre_dataset_transforms = dict()
post_dataset_transforms = dict()
for key in datasets_classes.keys():
pre_dataset_transforms[key] = transforms.Compose(pre_shared_transforms)
post_dataset_transforms[key] = transforms.Compose(
post_shared_transforms)
save_path = vae_util.setup_logging_from_args(args)
writer = SummaryWriter(save_path)
torch.manual_seed(args.seed)
cudnn.benchmark = True
torch.cuda.manual_seed(args.seed)
torch.cuda.manual_seed_all(args.seed)
if args.pred is not None:
checkpoint = torch.load(args.pred, map_location='cpu')
model_vae = model_vqvae.DecomposeNet(**checkpoint['arch_params'])
model_vae.load_state_dict(checkpoint['state_dict'])
else:
# FIXME: archs_param should be added to resume and fine_tune
arch_params = {'k': args.k, 'd': args.d, 'hidden': args.hidden}
model_vae = model_vqvae.DecomposeNet(hidden=args.hidden,
k=args.k,
d=args.d,
in_chns=args.in_chns,
out_chns=args.out_chns)
model_vae = model_vae.cuda()
# FIXME make it only for one single output
if args.lab_init:
distortion = [
116.0 / 500, 16.0 / 500, 500.0 / 500, 200.0 / 500, 0.2068966
]
trans_mat = [[0.412453, 0.357580, 0.180423],
[0.212671, 0.715160, 0.072169],
[0.019334, 0.119193, 0.950227]]
ref_white = (0.95047, 1., 1.08883)
tmat = colour_spaces.dkl_from_rgb.T
tmat = np.expand_dims(tmat, [2, 3])
cst_lr = args.lr * 0.1
else:
trans_mat = None
ref_white = None
distortion = None
tmat = None
cst_lr = args.lr
# model_cst = ColourTransformer.LabTransformer(
# trans_mat=trans_mat, ref_white=ref_white,
# distortion=distortion, linear=args.linear
# )
model_cst = ColourTransformer.ResNetTransformer(layers=args.cst_layers)
model_cst = model_cst.cuda()
vae_params = [
{
'params': [p for p in model_vae.parameters() if p.requires_grad]
},
]
cst_params = [
{
'params': [p for p in model_cst.parameters() if p.requires_grad]
},
]
optimizer_vae = optim.Adam(vae_params, lr=args.lr)
optimizer_cst = optim.Adam(cst_params, lr=cst_lr)
scheduler_vae = optim.lr_scheduler.StepLR(optimizer_vae,
int(args.epochs / 3), 0.5)
scheduler_cst = optim.lr_scheduler.StepLR(optimizer_cst,
int(args.epochs / 3), 0.5)
if args.resume is not None:
checkpoint = torch.load(args.resume, map_location='cpu')
model_vae.load_state_dict(checkpoint['state_dict'])
model_vae = model_vae.cuda()
args.start_epoch = checkpoint['epoch'] + 1
scheduler_vae.load_state_dict(checkpoint['scheduler_vae'])
optimizer_vae.load_state_dict(checkpoint['optimizer_vae'])
scheduler_cst.load_state_dict(checkpoint['scheduler_cst'])
optimizer_cst.load_state_dict(checkpoint['optimizer_cst'])
elif args.fine_tune is not None:
weights = torch.load(args.fine_tune, map_location='cpu')
model_vae.load_state_dict(weights, strict=False)
model_vae = model_vae.cuda()
intransform_funs = []
if args.in_space.lower() == 'cgi':
augmentation_settings = [{
'function': random_imutils.adjust_contrast,
'kwargs': {
'amount': np.array([0.2, 1.0]),
'channel_wise': True
}
}, {
'function': random_imutils.adjust_gamma,
'kwargs': {
'amount': np.array([0.2, 5.0]),
'channel_wise': True
}
}, {
'function': random_imutils.adjust_illuminant,
'kwargs': {
'illuminant': np.array([0.0, 1.0])
}
}]
intransform_funs.append(
cv2_preprocessing.RandomAugmentationTransformation(
augmentation_settings, num_augmentations=1))
elif args.in_space.lower() != 'rgb':
intransform_funs.append(
cv2_preprocessing.DecompositionTransformation(
args.in_space.lower()))
intransform = transforms.Compose(intransform_funs)
outtransform = None
args.outs_dict = {'rgb': {'vis_fun': None}}
# preparing the dataset
transforms_kwargs = {
'intransform': intransform,
'outtransform': outtransform,
'pre_transform': pre_dataset_transforms[args.dataset],
'post_transform': post_dataset_transforms[args.dataset]
}
if args.dataset in ['celeba', 'touch', 'ccvr']:
train_dataset = datasets_classes[args.dataset](root=args.data_dir,
split='train',
**transforms_kwargs)
test_dataset = datasets_classes[args.dataset](root=args.data_dir,
split='test',
**transforms_kwargs)
elif args.dataset in ['coco']:
train_dataset = datasets_classes[args.dataset](root=args.data_dir,
split='train',
**transforms_kwargs)
test_dataset = datasets_classes[args.dataset](root=args.data_dir,
split='val',
**transforms_kwargs)
elif args.dataset in ['voc']:
train_dataset = datasets_classes[args.dataset](root=args.data_dir,
image_set='train',
**transforms_kwargs)
test_dataset = datasets_classes[args.dataset](root=args.data_dir,
image_set='val',
**transforms_kwargs)
else:
train_dataset = datasets_classes[args.dataset](root=os.path.join(
args.data_dir, 'train'),
**transforms_kwargs)
test_dataset = datasets_classes[args.dataset](root=os.path.join(
args.data_dir, 'validation'),
**transforms_kwargs)
loader_kwargs = {
'batch_size': args.batch_size,
'num_workers': args.workers,
'pin_memory': True
}
train_loader = torch.utils.data.DataLoader(train_dataset,
shuffle=True,
**loader_kwargs)
test_loader = torch.utils.data.DataLoader(test_dataset,
shuffle=False,
**loader_kwargs)
if args.pred is not None:
predict(model_vae, test_loader, save_path, args)
return
# starting to train
console = logging.StreamHandler()
console.setLevel(logging.INFO)
logging.getLogger('').addHandler(console)
for epoch in range(args.start_epoch, args.epochs):
train_losses = train(epoch, model_vae, model_cst, train_loader,
(optimizer_vae, optimizer_cst), save_path, args)
test_losses = test_net(epoch, model_vae, model_cst, test_loader,
save_path, args)
for k in train_losses.keys():
name = k.replace('_trn', '')
train_name = k
test_name = k.replace('_trn', '_val')
writer.add_scalars(name, {
'train': train_losses[train_name],
'test': test_losses[test_name]
}, epoch)
scheduler_vae.step()
scheduler_cst.step()
vae_util.save_checkpoint(
{
'epoch': epoch,
'state_dict': model_vae.state_dict(),
'colour_transformer': model_cst.state_dict(),
'optimizer_vae': optimizer_vae.state_dict(),
'scheduler_vae': scheduler_vae.state_dict(),
'optimizer_cst': optimizer_cst.state_dict(),
'scheduler_cst': scheduler_cst.state_dict(),
'arch': args.model,
'arch_params': {
**arch_params,
'in_chns': args.in_chns,
'out_chns': args.out_chns,
},
'transformer_params': {
'linear': args.linear
}
}, save_path)
def main(args):
args = parse_arguments(args)
weights_rgb = torch.load(args.model_path, map_location='cpu')
network = vqmodel.VQ_CVAE(128,
k=args.k,
kl=args.kl,
in_chns=3,
cos_distance=args.cos_dis)
network.load_state_dict(weights_rgb)
if args.exclude > 0:
which_vec = [args.exclude - 1]
print(which_vec)
network.state_dict()['emb.weight'][:, which_vec] = 0
elif args.exclude < 0:
which_vec = [*range(8)]
which_vec.remove(abs(args.exclude) - 1)
print(which_vec)
network.state_dict()['emb.weight'][:, which_vec] = 0
network.cuda()
network.eval()
if not os.path.exists(args.out_dir):
os.mkdir(args.out_dir)
args.in_colour_space = args.colour_space[:3]
args.out_colour_space = args.colour_space[4:]
mean = (0.5, 0.5, 0.5)
std = (0.5, 0.5, 0.5)
transform_funcs = transforms.Compose([
# cv2_transforms.Resize(256), cv2_transforms.CenterCrop(224),
cv2_transforms.ToTensor(),
cv2_transforms.Normalize(mean, std)
])
intransform_funs = []
if args.in_colour_space != ' rgb':
intransform_funs.append(
cv2_preprocessing.ColourSpaceTransformation(args.in_colour_space))
intransform = transforms.Compose(intransform_funs)
if args.dataset == 'imagenet':
test_loader = torch.utils.data.DataLoader(data_loaders.ImageFolder(
root=args.validation_dir,
intransform=intransform,
outtransform=None,
transform=transform_funcs),
batch_size=args.batch_size,
shuffle=False)
else:
test_loader = torch.utils.data.DataLoader(
data_loaders.CategoryImages(
root=args.validation_dir,
# FIXME
category=args.category,
intransform=intransform,
outtransform=None,
transform=transform_funcs),
batch_size=args.batch_size,
shuffle=False)
export(test_loader, network, args)
def main(args):
args = arguments.parse_arguments(args)
# determining the number of input channels
in_size = 1
if args.in_space == 'gry':
args.in_chns = 1
elif args.in_space == 'db1':
args.in_chns = 4
in_size = 0.5
else:
args.in_chns = 3
# FIXME
# preparing the output dictionary
args.outs_dict = dict()
for out_type in args.outputs:
if out_type == 'input':
out_shape = [1 / in_size, 1 / in_size, args.in_chns]
elif out_type == 'gry':
out_shape = [1 / in_size, 1 / in_size, 1]
elif out_type == 'db1':
# TODO: just assuming numbers of square 2
out_shape = [0.5 / in_size, 0.5 / in_size, 4]
else:
out_shape = [1 / in_size, 1 / in_size, 3]
args.outs_dict[out_type] = {'shape': out_shape}
args.mean = 0.5
args.std = 0.5
target_size = args.target_size or dataset_target_size[args.dataset]
if args.dataset == 'ccvr':
pre_shared_transforms = [
cv2_transforms.Resize(target_size + 32),
cv2_transforms.RandomCrop(target_size),
]
else:
pre_shared_transforms = [
cv2_transforms.Resize(target_size + 32),
cv2_transforms.CenterCrop(target_size),
]
post_shared_transforms = [
cv2_transforms.ToTensor(),
cv2_transforms.Normalize(args.mean, args.std)
]
pre_dataset_transforms = dict()
post_dataset_transforms = dict()
for key in datasets_classes.keys():
pre_dataset_transforms[key] = transforms.Compose(pre_shared_transforms)
post_dataset_transforms[key] = transforms.Compose(
post_shared_transforms)
save_path = vae_util.setup_logging_from_args(args)
writer = SummaryWriter(save_path)
torch.manual_seed(args.seed)
cudnn.benchmark = True
torch.cuda.manual_seed(args.seed)
torch.cuda.manual_seed_all(args.seed)
if args.pred is not None:
checkpoint = torch.load(args.pred, map_location='cpu')
vae_model = (model_single
if checkpoint['model'] == 'single' else model_multi)
model = vae_model.DecomposeNet(**checkpoint['arch_params'])
model.load_state_dict(checkpoint['state_dict'])
else:
# FIXME: add to prediction, right now it's hard-coded for resnet18
if args.model == 'deeplabv3':
backbone = {
'arch': 'resnet_bottleneck_custom',
'customs': {
'pooling_type': 'max',
'in_chns': args.in_chns,
'blocks': [2, 2, 2, 2],
'num_kernels': 64,
'num_classes': 1000
}
}
# FIXME out_shape is defined far above, this is just a hack
arch_params = {'backbone': backbone, 'num_classes': out_shape[-1]}
model = model_segmentation.deeplabv3_resnet(
backbone, num_classes=out_shape[-1], outs_dict=args.outs_dict)
elif 'unet' in args.model:
from segmentation_models import unet
encoder_name = args.model.split('_')[-1]
model = unet.model.Unet(in_channels=args.in_chns,
encoder_name=encoder_name,
encoder_weights=None,
outs_dict=args.outs_dict,
classes=out_shape[-1])
arch_params = {'encoder_name': encoder_name}
elif args.model == 'category':
# FIXME: archs_param should be added to resume and fine_tune
arch_params = {'k': args.k, 'd': args.d, 'hidden': args.hidden}
vae_model = model_category
model = vae_model.DecomposeNet(hidden=args.hidden,
k=args.k,
d=args.d,
in_chns=args.in_chns,
outs_dict=args.outs_dict)
else:
# FIXME: archs_param should be added to resume and fine_tune
arch_params = {'k': args.k, 'd': args.d, 'hidden': args.hidden}
vae_model = model_single if args.model == 'single' else model_multi
model = vae_model.DecomposeNet(hidden=args.hidden,
k=args.k,
d=args.d,
in_chns=args.in_chns,
outs_dict=args.outs_dict)
model = model.cuda()
optimizer = optim.Adam(model.parameters(), lr=args.lr)
scheduler = optim.lr_scheduler.StepLR(optimizer, int(args.epochs / 3), 0.5)
if args.resume is not None:
checkpoint = torch.load(args.resume, map_location='cpu')
model.load_state_dict(checkpoint['state_dict'])
model = model.cuda()
args.start_epoch = checkpoint['epoch'] + 1
scheduler.load_state_dict(checkpoint['scheduler'])
optimizer.load_state_dict(checkpoint['optimizer'])
elif args.fine_tune is not None:
weights = torch.load(args.fine_tune, map_location='cpu')
model.load_state_dict(weights, strict=False)
model = model.cuda()
intransform_funs = []
if args.in_space.lower() == 'cgi':
augmentation_settings = [{
'function': random_imutils.adjust_contrast,
'kwargs': {
'amount': np.array([0.2, 1.0]),
'channel_wise': True
}
}, {
'function': random_imutils.adjust_gamma,
'kwargs': {
'amount': np.array([0.2, 5.0]),
'channel_wise': True
}
}, {
'function': random_imutils.adjust_illuminant,
'kwargs': {
'illuminant': np.array([0.0, 1.0])
}
}]
intransform_funs.append(
cv2_preprocessing.RandomAugmentationTransformation(
augmentation_settings, num_augmentations=1))
elif args.in_space.lower() != 'rgb':
intransform_funs.append(
cv2_preprocessing.DecompositionTransformation(
args.in_space.lower()))
intransform = transforms.Compose(intransform_funs)
outtransform_funs = [
cv2_preprocessing.MultipleOutputTransformation(args.outputs)
]
outtransform = transforms.Compose(outtransform_funs)
# FIXME
for out_type in args.outputs:
if out_type == 'input':
vis_fun = None
elif out_type == 'gry':
vis_fun = None
elif out_type == 'db1':
vis_fun = vae_util.wavelet_visualise
elif args.vis_rgb:
vis_fun = partial(all2rgb, src_space=out_type)
else:
vis_fun = None
args.outs_dict[out_type]['vis_fun'] = vis_fun
# preparing the dataset
transforms_kwargs = {
'intransform': intransform,
'outtransform': outtransform,
'pre_transform': pre_dataset_transforms[args.dataset],
'post_transform': post_dataset_transforms[args.dataset]
}
if args.dataset in ['celeba', 'touch', 'ccvr']:
train_dataset = datasets_classes[args.dataset](root=args.data_dir,
split='train',
**transforms_kwargs)
test_dataset = datasets_classes[args.dataset](root=args.data_dir,
split='test',
**transforms_kwargs)
elif args.dataset in ['coco']:
train_dataset = datasets_classes[args.dataset](root=args.data_dir,
split='train',
**transforms_kwargs)
test_dataset = datasets_classes[args.dataset](root=args.data_dir,
split='val',
**transforms_kwargs)
elif args.dataset in ['voc']:
train_dataset = datasets_classes[args.dataset](root=args.data_dir,
image_set='train',
**transforms_kwargs)
test_dataset = datasets_classes[args.dataset](root=args.data_dir,
image_set='val',
**transforms_kwargs)
else:
train_dataset = datasets_classes[args.dataset](root=os.path.join(
args.data_dir, 'train'),
**transforms_kwargs)
test_dataset = datasets_classes[args.dataset](root=os.path.join(
args.data_dir, 'validation'),
**transforms_kwargs)
loader_kwargs = {
'batch_size': args.batch_size,
'num_workers': args.workers,
'pin_memory': True
}
train_loader = torch.utils.data.DataLoader(train_dataset,
shuffle=True,
**loader_kwargs)
test_loader = torch.utils.data.DataLoader(test_dataset,
shuffle=False,
**loader_kwargs)
if args.pred is not None:
predict(model, test_loader, save_path, args)
return
# starting to train
console = logging.StreamHandler()
console.setLevel(logging.INFO)
logging.getLogger('').addHandler(console)
for epoch in range(args.start_epoch, args.epochs):
train_losses = train(epoch, model, train_loader, optimizer, save_path,
args)
test_losses = test_net(epoch, model, test_loader, save_path, args)
for k in train_losses.keys():
name = k.replace('_train', '')
train_name = k
test_name = k.replace('train', 'test')
writer.add_scalars(name, {
'train': train_losses[train_name],
'test': test_losses[test_name]
}, epoch)
scheduler.step()
vae_util.save_checkpoint(
{
'epoch': epoch,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict(),
'scheduler': scheduler.state_dict(),
'arch': args.model,
'arch_params': {
**arch_params, 'in_chns': args.in_chns,
'outs_dict': args.outs_dict
}
}, save_path)
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!')
layer_name = args.transfer_weights[0]
weights_path = None
if len(args.transfer_weights) > 1:
weights_path = args.transfer_weights[1]
model = contrast_utils.AFCModel(args.network_name,
layer_name=layer_name,
weights_path=weights_path,
num_classes=args.num_classes)
# FIXME: remove me
# from kernelphysiology.dl.pytorch import models as custom_models
# model = custom_models.__dict__[args.network_name](
# in_chns=len(mean), num_classes=args.num_classes
# )
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
criterion = nn.MSELoss().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:
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.AADB(root=args.data_dir,
split='train',
transform=train_trans)
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.AADB(root=args.data_dir,
split='test',
transform=valid_trans)
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': {
'layer_name': layer_name,
'in_chns': len(mean),
'num_classes': args.num_classes,
'weights_path': weights_path,
},
'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,v_time,v_loss,v_top1'
np.savetxt(model_progress_path,
np.array(model_progress),
delimiter=',',
header=header)
