def __getitem__(self, index):
if self.is_train:
img, target = imageio.imread(
self.train_img_label[index][0]), self.train_img_label[index][1]
if len(img.shape) == 2:
img = np.stack([img] * 3, 2)
img = Image.fromarray(img, mode='RGB')
img = transforms.Resize((self.input_size, self.input_size),
Image.BILINEAR)(img)
if (rand_aug):
img = RandAugment(N, M)(img)
# img = transforms.RandomResizedCrop(size=self.input_size,scale=(0.4, 0.75),ratio=(0.5,1.5))(img)
# img = transforms.RandomCrop(self.input_size)(img)
img = transforms.RandomHorizontalFlip()(img)
img = transforms.ColorJitter(brightness=0.2, contrast=0.2)(img)
img = transforms.ToTensor()(img)
img = transforms.Normalize([0.485, 0.456, 0.406],
[0.229, 0.224, 0.225])(img)
else:
img, target = imageio.imread(
self.test_img_label[index][0]), self.test_img_label[index][1]
if len(img.shape) == 2:
img = np.stack([img] * 3, 2)
img = Image.fromarray(img, mode='RGB')
img = transforms.Resize((self.input_size, self.input_size),
Image.BILINEAR)(img)
# img = transforms.CenterCrop(self.input_size)(img)
img = transforms.ToTensor()(img)
img = transforms.Normalize([0.485, 0.456, 0.406],
[0.229, 0.224, 0.225])(img)
return img, target
def randagument(self):
self.transform = transforms.Compose([
transforms.CenterCrop(self.size),
RandAugment(1, 15),
ToTensor(),
Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
])
def UnlabeledImageDataLoader(transform_args,
randaug = False,
image_dir = '../../ChestX-ray14/images/',
label_dir_list = ['../../ChestX-ray14/train_val_list.txt', '../../ChestX-ray14/test_list.txt'],
batchsize = 64,
numworker = 12,
res = 456,
train = True,
):
# --- Transforms ---
transform_list = [getattr(transforms, transform_args[0])(res),
*[getattr(transforms, single_transform)() for single_transform in transform_args[1:]]]
img_transform = transforms.Compose(transform_list)
if train and randaug:
img_transform.transforms.insert(0, RandAugment(2, 3))
# --- Data Collection ---
unlabeled_dataset = UnlabeledImageDataset(image_dir=image_dir,
label_dir_list=label_dir_list,
transform=img_transform)
# --- DataLoader ---
unlabeled_loader = torch.utils.data.DataLoader(unlabeled_dataset, batch_size=batchsize,
num_workers=numworker, pin_memory=True)
return unlabeled_loader, unlabeled_dataset.label_dir
def randaugment(self):
self.transform = transforms.Compose([
transforms.CenterCrop(self.size),
transforms.RandomHorizontalFlip(p=0.5),
RandAugment(1, 15),
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
])
def get_simclr_data_transforms_randAugment(input_shape):
# get a set of data augmentation transformations as described in the SimCLR paper.
data_transforms = transforms.Compose([
transforms.RandomResizedCrop(size=eval(input_shape)[0]),
transforms.RandomHorizontalFlip(),
transforms.ToTensor()
])
data_transforms.transforms.insert(0, RandAugment(3, 9))
# pip install git+https://github.com/ildoonet/pytorch-randaugment
return data_transforms
def __init__(self, name, data, targets, transforms, base_transforms): self.name = name self.data = data self.targets = targets self.transforms = transforms self.base_transforms = base_transforms self.rand_aug_transforms = copy.deepcopy(self.base_transforms) self.committee_size = 1 self.ra_obj = RandAugment(1, 2.0) if self.name == 'mnist': # See if this can be dropped entirely without losing performance self.transforms.transforms.insert(0, torchvision.transforms.Grayscale(num_output_channels=1)) self.base_transforms.transforms.insert(0, torchvision.transforms.Grayscale(num_output_channels=1)) self.rand_aug_transforms.transforms.insert(0, torchvision.transforms.Grayscale(num_output_channels=1)) self.rand_aug_transforms.transforms.insert(0, self.ra_obj)
def prepare_data(self):
##### RandAugment changes here for now
data_transforms = { # same for now
'train': transforms.Compose([
transforms.Resize((256, 256), interpolation=2),
# transforms.RandomCrop(32, padding=0),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(mean=[n / 255.
for n in
[129.3, 124.1,
112.4]],
std=[n / 255. for n in
[68.2, 65.4,
70.4]])
]),
'valid': transforms.Compose([
transforms.Resize((256, 256), interpolation=2),
# transforms.RandomCrop(32, padding=0),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(mean=[n / 255.
for n in
[129.3, 124.1,
112.4]],
std=[n / 255. for n in
[68.2, 65.4,
70.4]])
]),
}
data_transforms['train'].transforms.insert(
0, RandAugment(self.randAugParamN, self.randAugParamM))
train_path = '/lab/vislab/OPEN/datasets_RGB_new/train/'
valid_path = '/lab/vislab/OPEN/datasets_RGB_new/val/'
self.trainset = torchvision.datasets.ImageFolder(
train_path, data_transforms['train'])
self.validset = torchvision.datasets.ImageFolder(
valid_path, data_transforms['valid'])
def ImageDataLoader(transform_args,
pseudo_label_args,
randaug = False,
image_dir = '../../',
label_dir = '../../CheXpert-v1.0-small/train.csv',
batchsize = 64,
numworker = 12,
res = 456,
train = True,
label_smooth = False,
):
"""
train_dataset (images):
[0]: Minibatch of Images. torch.Tensor of size (batchsize, 3, H, W)
[1]: What subfolder the data come from. list of size (batchsize)
train_label:
1: Positive
0: Negative + Unmentioned
-1: Uncertain
"""
# --- Transforms ---
transform_list = [getattr(transforms, transform_args[0])(res),
*[getattr(transforms, single_transform)() for single_transform in transform_args[1:]]]
img_transform = transforms.Compose(transform_list)
if train and randaug:
img_transform.transforms.insert(0, RandAugment(2, 3))
# --- Data Collection ---
train_dataset = ImageDataset(image_dir=image_dir,
label_dir=label_dir,
transform=img_transform,
train=train,
label_smooth=label_smooth)
if pseudo_label_args['use_pseudo_label']:
ssl_dataset = ImageDataset(image_dir=pseudo_label_args['pseudo_dataset_args']['image_dir'],
label_dir=pseudo_label_args['pseudo_dataset_args']['label_dir'],
transform=img_transform,
train=train,
label_smooth=label_smooth)
train_dataset = torch.utils.data.ConcatDataset((train_dataset, ssl_dataset))
# --- DataLoader ---
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batchsize, num_workers=numworker, pin_memory=True)
return train_loader, len(train_dataset)
def load_dataset(N, M, dataset, dataset_dir):
'''Add RandAugment with N, M(hyperparameter)
N: Number of augmentation transformations to apply sequentially.
M: Magnitude for all the transformations.'''
transform_randaugment = transforms.Compose([
RandAugment(N, M),
#transforms.Resize(32),
transforms.ToTensor(),
])
transform_resize = transforms.Compose([
#transforms.Resize(32),
transforms.ToTensor(),
])
if dataset == 'stl10':
stl10 = {}
stl10['train_augmented'] = STL10(dataset_dir,
'train',
transform=transform_randaugment,
target_transform=one_hot,
download=True)
stl10['train_unaugmented'] = STL10(dataset_dir,
'train',
transform=transform_resize,
target_transform=one_hot)
stl10['unlabeled_augmented'] = STL10(dataset_dir,
'unlabeled',
transform=transform_randaugment,
download=True)
stl10['unlabeled_unaugmented'] = STL10(dataset_dir,
'unlabeled',
transform=transform_resize,
download=True)
stl10['test'] = STL10(dataset_dir,
'test',
transform=transform_resize,
download=True)
return stl10
else:
raise Excpetion(f"Dataset '{dataset}' not implemented")
def get_dataset(dataset, cutout_length=0, N=3, M=5, RandA=False):
MEAN = [0.49139968, 0.48215827, 0.44653124]
STD = [0.24703233, 0.24348505, 0.26158768]
transf = [
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip()
]
normalize = [
transforms.Resize(32),
transforms.ToTensor(),
transforms.Normalize(MEAN, STD)
]
cutout = []
if cutout_length > 0:
cutout.append(Cutout(cutout_length))
train_transform = transforms.Compose(transf + normalize + cutout)
if RandA:
# Add RandAugment with N, M(hyperparameter)
train_transform.transforms.insert(0, RandAugment(N, M))
valid_transform = transforms.Compose(normalize)
if dataset == "cifar10":
dataset_train = CIFAR10(root="./data", train=True, download=True, transform=train_transform)
dataset_valid = CIFAR10(root="./data", train=False, download=True, transform=valid_transform)
elif dataset == 'cifar100':
dataset_train = CIFAR100(root="./data", train=True, download=True, transform=train_transform)
dataset_valid = CIFAR100(root="./data", train=False, download=True, transform=valid_transform)
else:
raise NotImplementedError
return dataset_train, dataset_valid
import argparse
def parse_args():
parser = argparse.ArgumentParser(
description = "Arguments for ModelNet Visualization",
formatter_class = argparse.ArgumentDefaultsHelpFormatter,
)
parser.add_argument("-N", type=int, default=4, help="RandAugment N")
parser.add_argument("-M", type=int, default=4, help="RandAugment M")
return parser.parse_args()
args = parse_args()
transforms = transforms.Compose(
[
RandAugment(n=args.N, m=args.M),
d_utils.PointCloudToTensor(),
d_utils.PointCloudNormalize(),
]
)
dset = ModelNet40Cls(1024, train = True, transforms=transforms)
pc , _ = dset[100]
pc = pc[:,0:3]
#dloader = torch.utils.data.DataLoader(dset, batch_size = 1, shuffle=True)
xyz = pc.numpy()
# Pass xyz to Open3D.o3d.geometry.PointCloud and visualize
pcd = o3d.geometry.PointCloud()
pcd.points = o3d.utility.Vector3dVector(xyz)
o3d.io.write_point_cloud("sync.ply", pcd)
def main_worker(gpu, ngpus_per_node, args):
CHECKPOINT_ID = "{}_{}epochs_{}bsz_{:0.4f}lr" \
.format(args.id[:5], args.epochs, args.batch_size, args.lr)
if args.mlp:
CHECKPOINT_ID += "_mlp"
if args.aug_plus:
CHECKPOINT_ID += "_augplus"
if args.cos:
CHECKPOINT_ID += "_cos"
if args.faa_aug:
CHECKPOINT_ID += "_faa"
if args.randomcrop:
CHECKPOINT_ID += "_randcrop"
if args.rotnet:
CHECKPOINT_ID += "_rotnet"
if args.rand_aug:
CHECKPOINT_ID += "_randaug"
if not (args.kfold == None):
CHECKPOINT_ID += "_fold_%d" % (args.kfold)
if not (args.custom_aug_name == None):
CHECKPOINT_ID += "_custom_aug_" + args.custom_aug_name
CHECKPOINT_ID += args.dataid
args.gpu = gpu
# suppress printing if not master
if args.multiprocessing_distributed and args.gpu != 0:
def print_pass(*args):
pass
builtins.print = print_pass
if args.gpu is not None:
print("Use GPU: {} for training".format(args.gpu))
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 + gpu
dist.init_process_group(backend=args.dist_backend,
init_method=args.dist_url,
world_size=args.world_size,
rank=args.rank)
# create model
heads = {}
if not args.nomoco:
heads["moco"] = {"num_classes": args.moco_dim}
if args.rotnet:
heads["rotnet"] = {"num_classes": 4}
model = moco.builder.MoCo(models.__dict__[args.arch],
K=args.moco_k,
m=args.moco_m,
T=args.moco_t,
mlp=args.mlp,
dataid=args.dataid,
multitask_heads=heads)
print(model)
# setup file structure for saving
pathlib.Path(args.checkpoint_fp).mkdir(parents=True, exist_ok=True)
if args.distributed:
# For multiprocessing distributed, DistributedDataParallel constructor
# should always set the single device scope, otherwise,
# DistributedDataParallel will use all available devices.
if args.gpu is not None:
torch.cuda.set_device(args.gpu)
model.cuda(args.gpu)
# 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 - 1) / ngpus_per_node)
model = torch.nn.parallel.DistributedDataParallel(
model, device_ids=[args.gpu])
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.gpu is not None:
torch.cuda.set_device(args.gpu)
model = model.cuda(args.gpu)
# comment out the following line for debugging
# raise NotImplementedError("Only DistributedDataParallel is supported.")
else:
# AllGather implementation (batch shuffle, queue update, etc.) in
# this code only supports DistributedDataParallel.
raise NotImplementedError("Only DistributedDataParallel is supported.")
# define loss function (criterion) and optimizer
criterion = nn.CrossEntropyLoss().cuda(args.gpu)
optimizer = torch.optim.SGD(model.parameters(),
args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
# optionally resume from a checkpoint
if args.resume:
if os.path.isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
if args.gpu is None:
checkpoint = torch.load(args.resume)
else:
# Map model to be loaded to specified single gpu.
loc = 'cuda:{}'.format(args.gpu)
checkpoint = torch.load(args.resume, map_location=loc)
args.start_epoch = checkpoint['epoch']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
args.id = checkpoint['id']
args.name = checkpoint['name']
CHECKPOINT_ID = checkpoint['name']
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(args.resume))
cudnn.benchmark = True
# Data loading code
# Set up crops and normalization depending on the dataset.
# Cifar 10 crops and normalization.
if args.dataid == "cifar10" or args.dataid == "svhn":
_CIFAR_MEAN, _CIFAR_STD = (0.4914, 0.4822, 0.4465), (0.2023, 0.1994,
0.2010)
normalize = transforms.Normalize(mean=_CIFAR_MEAN, std=_CIFAR_STD)
if not args.randomcrop:
random_resized_crop = transforms.RandomResizedCrop(
28, scale=(args.rrc_param, 1.))
else:
# Use the crop they were using in Fast AutoAugment.
random_resized_crop = transforms.RandomCrop(32, padding=4)
# Use the imagenet parameters.
elif args.dataid == "imagenet" or args.dataid == "logos":
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
random_resized_crop = transforms.RandomResizedCrop(224,
scale=(0.2, 1.))
if args.aug_plus and (args.faa_aug or args.rand_aug or args.rand_aug_orig
or not (args.custom_aug_name == None)):
raise Exception("Cannot have multiple augs on command line")
if args.aug_plus:
# MoCo v2's aug: similar to SimCLR https://arxiv.org/abs/2002.05709
augmentation = [
random_resized_crop,
transforms.RandomApply(
[
transforms.ColorJitter(0.4, 0.4, 0.4,
0.1) # not strengthened
],
p=0.8),
transforms.RandomGrayscale(p=0.2),
transforms.RandomApply(
[moco.loader.GaussianBlur([args.sigma / 20, args.sigma])],
p=0.5),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize
]
elif args.faa_aug:
augmentation, _ = slm_utils.get_faa_transforms.get_faa_transforms_cifar_10(
args.randomcrop, args.gauss)
transformations = moco.loader.TwoCropsTransform(augmentation)
elif args.rand_aug_orig:
print("Using random aug original")
augmentation = [
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
RandAugment(args.rand_aug_n, args.rand_aug_m),
transforms.ToTensor(), normalize
]
elif args.rand_aug:
randaug_n = args.rand_aug_n
if args.rand_aug_linear_m:
print("Using random aug with linear m")
randaug_m = args.rand_aug_m_min
else:
randaug_m = args.rand_aug_m
print("Using random aug")
if args.rand_aug_top_k > 0:
randaug = TopRandAugment(randaug_n, randaug_m, args.rand_aug_top_k)
else:
randaug = RandAugment(randaug_n, randaug_m)
augmentation = [
random_resized_crop,
transforms.RandomHorizontalFlip(), randaug,
transforms.ToTensor(), normalize
]
elif args.rand_resize_only and args.custom_aug_name == None:
print("Using random resize only")
augmentation = [
random_resized_crop,
transforms.RandomHorizontalFlip(),
transforms.ToTensor(), normalize
]
elif not args.custom_aug_name == None:
augmentation, _ = slm_utils.get_faa_transforms.load_custom_transforms(
name=args.custom_aug_name,
randomcrop=args.randomcrop,
aug_idx=args.single_aug_idx,
dataid=args.dataid)
print('using custom augs', augmentation)
transformations = moco.loader.TwoCropsTransform(augmentation)
else:
# MoCo v1's aug: the same as InstDisc https://arxiv.org/abs/1805.01978
print('using v1 augs')
augmentation = [
random_resized_crop,
transforms.RandomGrayscale(p=0.2),
transforms.ColorJitter(0.4, 0.4, 0.4, 0.4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(), normalize
]
if not args.faa_aug and args.custom_aug_name == None:
print('using augmentation', augmentation)
transformations = moco.loader.TwoCropsTransform(
transforms.Compose(augmentation))
print('xforms', transformations)
if args.dataid == "imagenet" and not args.reduced_imgnet:
train_dataset = datasets.ImageFolder(args.data, transformations)
elif args.dataid == "logos" and not args.reduced_imgnet:
train_dataset = data_loader.GetLoader(
data_root=args.data,
data_list='train_images_root.txt',
transform=transformations)
elif (args.dataid == "imagenet"
or args.dataid == 'logos') and args.reduced_imgnet:
# idx120 = [16, 23, 52, 57, 76, 93, 95, 96, 99, 121, 122, 128, 148, 172, 181, 189, 202, 210, 232, 238, 257, 258, 259, 277, 283, 289, 295, 304, 307, 318, 322, 331, 337, 338, 345, 350, 361, 375, 376, 381, 388, 399, 401, 408, 424, 431, 432, 440, 447, 462, 464, 472, 483, 497, 506, 512, 530, 541, 553, 554, 557, 564, 570, 584, 612, 614, 619, 626, 631, 632, 650, 657, 658, 660, 674, 675, 680, 682, 691, 695, 699, 711, 734, 736, 741, 754, 757, 764, 769, 770, 780, 781, 787, 797, 799, 811, 822, 829, 830, 835, 837, 842, 843, 845, 873, 883, 897, 900, 902, 905, 913, 920, 925, 937, 938, 940, 941, 944, 949, 959]
if args.dataid == "imagenet":
total_trainset = ImageNet(
root=args.data, transform=transformations
) # TODO for LINCLS, make this train and test xforms.
else:
total_trainset = data_loader.GetLoader(
data_root=args.data,
data_list='train_images_root.txt',
transform=transformations)
train_idx = np.arange(len(total_trainset))
np.random.seed(1337) #fingers crossed.
np.random.shuffle(train_idx)
train_idx = train_idx[:50000]
kfold = args.kfold
print('KFOLD BEING USED', kfold)
subset = np.arange(kfold * 10000, (kfold + 1) * 10000)
print('start', 'end', kfold * 10000, (kfold + 1) * 10000)
valid_idx = train_idx[subset]
train_idx = np.delete(train_idx, subset)
print('first val_idx', valid_idx[:10])
train_dataset = total_trainset
train_dataset = Subset(train_dataset, train_idx)
train_sampler = SubsetRandomSampler(train_idx)
valid_sampler = SubsetSampler(valid_idx)
print(len(train_dataset))
train_sampler = torch.utils.data.distributed.DistributedSampler(
train_dataset)
print(len(train_dataset))
print('first 10 train', train_idx[:10])
print('first 10 valid', valid_idx[:10])
print('len train', len(train_idx))
print('len valid', len(valid_idx))
for i in valid_idx:
if i in train_idx:
raise Exception("Valid idx in train idx: this is unexpected")
print('train_sampler', train_sampler)
elif args.dataid == "cifar10":
train_dataset = torchvision.datasets.CIFAR10(args.data,
transform=transformations,
download=True)
elif args.dataid == "svhn":
train_dataset = torchvision.datasets.SVHN(args.data,
transform=transformations,
download=True)
else:
raise NotImplementedError(
"Support for the following dataset is not yet implemented: {}".
format(args.dataid))
if not args.kfold == None and not args.reduced_imgnet:
torch.manual_seed(1337)
print('before: K FOLD', args.kfold, len(train_dataset))
lengths = [len(train_dataset) // 5] * 5
print(lengths)
lengths[-1] = int(lengths[-1] + (len(train_dataset) - np.sum(lengths)))
print(lengths)
folds = torch.utils.data.random_split(train_dataset, lengths)
print(len(folds))
folds.pop(args.kfold)
print(len(folds))
train_dataset = torch.utils.data.ConcatDataset(folds)
print(len(train_dataset))
else:
print("NO KFOLD ARG", args.kfold, ' or ', args.reduced_imgnet)
if args.distributed and not args.reduced_imgnet:
train_sampler = torch.utils.data.distributed.DistributedSampler(
train_dataset)
elif not args.reduced_imgnet:
train_sampler = None
print('train sampler', train_sampler)
torch.manual_seed(1337)
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,
drop_last=True)
print(len(train_loader))
# CR: only the master will report to wandb for now
if not args.multiprocessing_distributed or args.rank % ngpus_per_node == 0:
wandb.init(project=args.wandbproj,
name=CHECKPOINT_ID,
id=args.id,
resume=args.resume,
config=args.__dict__,
notes=args.notes)
print(model)
for epoch in range(args.start_epoch, args.epochs):
if args.distributed:
train_sampler.set_epoch(epoch)
adjust_learning_rate(optimizer, epoch, args)
if args.rand_aug_linear_m:
mval = args.rand_aug_m_min + math.floor(
float(epoch) / float(args.epochs) *
(args.rand_aug_m_max - args.rand_aug_m_min + 1))
print("Rand aug m: {}".format(mval))
randaug.m = mval
# train for one epoch
train(train_loader, model, criterion, optimizer, epoch, args,
CHECKPOINT_ID)
# save current epoch
if not args.nosave_latest and (not args.multiprocessing_distributed
or args.rank % ngpus_per_node == 0):
print("saving latest epoch")
cp_filename = "{}_latest.tar".format(CHECKPOINT_ID[:5])
cp_fullpath = os.path.join(args.checkpoint_fp, cp_filename)
torch.save(
{
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict(),
'id': args.id,
'name': CHECKPOINT_ID,
}, cp_fullpath)
print("saved latest epoch")
if (epoch % args.checkpoint_interval == 0 or epoch == args.epochs-1) \
and (not args.multiprocessing_distributed or
(args.multiprocessing_distributed and args.rank % ngpus_per_node == 0)):
cp_filename = "{}_{:04d}.tar".format(CHECKPOINT_ID, epoch)
cp_fullpath = os.path.join(args.checkpoint_fp, cp_filename)
torch.save(
{
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict(),
'id': args.id,
'name': CHECKPOINT_ID,
}, cp_fullpath)
if args.upload_checkpoints:
print("Uploading wandb checkpoint")
wandb.save(cp_fullpath)
if epoch == args.epochs - 1:
print("Saving final results to wandb")
wandb.save(cp_fullpath)
print("Done - wrapping up")
def configure_transform(phase: str, transform_type: str):
if phase == "train":
if transform_type == Aug.BaseTransform:
transform = transforms.Compose([
transforms.Resize((512, 384), Image.BILINEAR),
transforms.CenterCrop((384, 384)),
transforms.RandomResizedCrop((224, 224)),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406],
[0.229, 0.224, 0.225]),
])
elif transform_type == Aug.FaceCrop:
transform = transforms.Compose([
transforms.Lambda(lambda x: crop(x)),
transforms.Resize((312, 234)),
transforms.RandomCrop((224, 224)),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406],
[0.229, 0.224, 0.225]),
])
elif transform_type == Aug.Random:
transform = transforms.Compose([
transforms.CenterCrop((384, 384)),
transforms.RandomResizedCrop((224, 224)),
RandAugment(2, 9), # N: 몇 개 선택할지 M: 몇 번 변화시킬 것인지
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406],
[0.229, 0.224, 0.225]),
])
elif transform_type == Aug.TTA: # 'random'과 동일
transform = transforms.Compose([
transforms.CenterCrop((384, 384)),
transforms.RandomResizedCrop((224, 224)),
RandAugment(2, 9), # N: 몇 개 선택할지 M: 몇 번 변화시킬 것인지
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406],
[0.229, 0.224, 0.225]),
])
else:
raise NotImplementedError()
else:
if transform_type == Aug.BaseTransform:
transform = transforms.Compose([
transforms.Resize((512, 384), Image.BILINEAR),
transforms.CenterCrop((224, 224)),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406],
[0.229, 0.224, 0.225]),
])
elif transform_type == Aug.FaceCrop:
transform = transforms.Compose([
transforms.Lambda(lambda x: crop(x)),
transforms.Resize((312, 234)),
transforms.CenterCrop((224, 224)),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406],
[0.229, 0.224, 0.225]),
])
elif transform_type == Aug.Random:
transform = transforms.Compose([
transforms.CenterCrop((224, 224)),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406],
[0.229, 0.224, 0.225]),
])
elif transform_type == Aug.TTA:
transform = transforms.Compose([
transforms.CenterCrop((384, 384)),
transforms.RandomResizedCrop((224, 224)),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406],
[0.229, 0.224, 0.225]),
])
else:
raise NotImplementedError()
return transform
# Augmentations.
transform_train = transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(cifar10_mean, cifar10_std)
])
transform_strong = transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(cifar10_mean, cifar10_std)
])
transform_strong.transforms.insert(0, RandAugment(3, 4))
transform_strong.transforms.append(CutoutDefault(16))
transform_val = transforms.Compose(
[transforms.ToTensor(),
transforms.Normalize(cifar10_mean, cifar10_std)])
class TransformMix:
def __init__(self, transform):
self.transform = transform
def __call__(self, inp):
out1 = self.transform(inp)
out2 = self.transform(inp)
return out1, out2
def main():
args = parse_args()
if args.name is None:
args.name = '%s_%s' % (args.arch, datetime.now().strftime('%m%d%H%M'))
if not os.path.exists('models/%s' % args.name):
os.makedirs('models/%s' % args.name)
if args.print_model:
print('Config -----')
for arg in vars(args):
print('- %s: %s' % (arg, getattr(args, arg)))
print('------------')
with open('models/%s/args.txt' % args.name, 'w') as f:
for arg in vars(args):
print('- %s: %s' % (arg, getattr(args, arg)), file=f)
joblib.dump(args, 'models/%s/args.pkl' % args.name)
if args.loss == 'CrossEntropyLoss':
criterion = nn.CrossEntropyLoss().cuda()
elif args.loss == 'FocalLoss':
criterion = FocalLoss().cuda()
elif args.loss == 'MSELoss':
criterion = nn.MSELoss().cuda()
elif args.loss == 'BCELoss':
criterion = nn.BCELoss().cuda()
elif args.loss == 'L1Loss':
criterion = nn.L1Loss().cuda()
elif args.loss == 'multitask':
criterion = {
'classification': nn.CrossEntropyLoss().cuda(),
'regression': nn.MSELoss().cuda(),
}
else:
raise NotImplementedError
if args.pred_type == 'classification':
num_outputs = 5
elif args.pred_type == 'regression':
num_outputs = 1
elif args.loss == 'multitask':
num_outputs = 6
else:
raise NotImplementedError
cudnn.benchmark = True
train_transform = []
if args.RAugment:
train_transform = transforms.Compose([
transforms.Resize((args.img_size, args.img_size)),
transforms.RandomAffine(
degrees=(args.rotate_min,
args.rotate_max) if args.rotate else 0,
translate=(args.translate_min,
args.translate_max) if args.translate else None,
scale=(args.rescale_min,
args.rescale_max) if args.rescale else None,
shear=(args.shear_min, args.shear_max) if args.shear else None,
),
transforms.CenterCrop(args.input_size),
transforms.RandomHorizontalFlip(p=0.5 if args.flip else 0),
# transforms.RandomVerticalFlip(p=0.5 if args.flip else 0),
transforms.ColorJitter(
brightness=(args.brightness_min,
args.brightness_max) if args.brightness else 0,
contrast=(args.contrast_min,
args.contrast_max) if args.contrast else 0,
saturation=(args.saturation_min,
args.saturation_max) if args.saturation else 0,
hue=0),
RandomErase(
prob=args.random_erase_prob if args.random_erase else 0,
sl=args.random_erase_sl,
sh=args.random_erase_sh,
r=args.random_erase_r),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])
train_transform.transforms.insert(
0, RandAugment(args.RAugment_N, args.RAugment_M))
else:
train_transform = transforms.Compose([
transforms.Resize((args.img_size, args.img_size)),
transforms.RandomAffine(
degrees=(args.rotate_min,
args.rotate_max) if args.rotate else 0,
translate=(args.translate_min,
args.translate_max) if args.translate else None,
scale=(args.rescale_min,
args.rescale_max) if args.rescale else None,
shear=(args.shear_min, args.shear_max) if args.shear else None,
),
transforms.CenterCrop(args.input_size),
transforms.RandomHorizontalFlip(p=0.5 if args.flip else 0),
# transforms.RandomVerticalFlip(p=0.5 if args.flip else 0),
transforms.ColorJitter(
brightness=(args.brightness_min,
args.brightness_max) if args.brightness else 0,
contrast=(args.contrast_min,
args.contrast_max) if args.contrast else 0,
saturation=(args.saturation_min,
args.saturation_max) if args.saturation else 0,
hue=0),
RandomErase(
prob=args.random_erase_prob if args.random_erase else 0,
sl=args.random_erase_sl,
sh=args.random_erase_sh,
r=args.random_erase_r),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])
# IF Class Aware, need to do Data Augumentation on Validation.
if args.class_aware:
val_transform = transforms.Compose([
transforms.Resize((args.img_size, args.input_size)),
transforms.RandomAffine(
degrees=(args.rotate_min,
args.rotate_max) if args.rotate else 0,
translate=(args.translate_min,
args.translate_max) if args.translate else None,
scale=(args.rescale_min,
args.rescale_max) if args.rescale else None,
shear=(args.shear_min, args.shear_max) if args.shear else None,
),
transforms.CenterCrop(args.input_size),
transforms.RandomHorizontalFlip(p=0.5 if args.flip else 0),
# transforms.RandomVerticalFlip(p=0.5 if args.flip else 0),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])
else:
val_transform = transforms.Compose([
transforms.Resize((args.img_size, args.input_size)),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])
# Data loading.
if 'diabetic_retinopathy' in args.train_dataset:
diabetic_retinopathy_dir = preprocess('diabetic_retinopathy',
args.img_size,
scale=args.scale_radius,
norm=args.normalize,
pad=args.padding,
remove=args.remove)
diabetic_retinopathy_df = pd.read_csv(
'inputs/diabetic-retinopathy-resized/trainLabels.csv')
diabetic_retinopathy_img_paths = \
diabetic_retinopathy_dir + '/' + diabetic_retinopathy_df['image'].values + '.jpeg'
diabetic_retinopathy_labels = diabetic_retinopathy_df['level'].values
if 'RNFL' in args.train_dataset:
aptos2019_dir = preprocess('RNFL',
args.img_size,
scale=args.scale_radius,
norm=args.normalize,
pad=args.padding,
remove=args.remove,
red_free=args.red_free)
aptos2019_df = pd.read_csv(Path("./data/dataset/train/label.csv"))
aptos2019_img_paths = aptos2019_dir + '/' + aptos2019_df[
'id_code'].values + '.png'
aptos2019_labels = aptos2019_df['diagnosis'].values
if args.train_dataset == 'RNFL':
skf = StratifiedKFold(n_splits=args.n_splits,
shuffle=True,
random_state=41)
img_paths = []
labels = []
for fold, (train_idx, val_idx) in enumerate(
skf.split(aptos2019_img_paths, aptos2019_labels)):
img_paths.append(
(aptos2019_img_paths[train_idx], aptos2019_img_paths[val_idx]))
labels.append(
(aptos2019_labels[train_idx], aptos2019_labels[val_idx]))
elif args.train_dataset == 'diabetic_retinopathy':
img_paths = [(diabetic_retinopathy_img_paths, aptos2019_img_paths)]
labels = [(diabetic_retinopathy_labels, aptos2019_labels)]
elif 'diabetic_retinopathy' in args.train_dataset and 'aptos2019' in args.train_dataset:
skf = StratifiedKFold(n_splits=args.n_splits,
shuffle=True,
random_state=41)
img_paths = []
labels = []
for fold, (train_idx, val_idx) in enumerate(
skf.split(aptos2019_img_paths, aptos2019_labels)):
img_paths.append((np.hstack((aptos2019_img_paths[train_idx],
diabetic_retinopathy_img_paths)),
aptos2019_img_paths[val_idx]))
labels.append((np.hstack(
(aptos2019_labels[train_idx], diabetic_retinopathy_labels)),
aptos2019_labels[val_idx]))
# else:
# raise NotImplementedError
if args.pseudo_labels:
test_df = pd.read_csv('probs/%s.csv' % args.pseudo_labels)
test_dir = preprocess('test',
args.img_size,
scale=args.scale_radius,
norm=args.normalize,
pad=args.padding,
remove=args.remove)
test_img_paths = test_dir + '/' + test_df['id_code'].values + '.png'
test_labels = test_df['diagnosis'].values
for fold in range(len(img_paths)):
img_paths[fold] = (np.hstack(
(img_paths[fold][0], test_img_paths)), img_paths[fold][1])
labels[fold] = (np.hstack(
(labels[fold][0], test_labels)), labels[fold][1])
if 'messidor' in args.train_dataset:
test_dir = preprocess('messidor',
args.img_size,
scale=args.scale_radius,
norm=args.normalize,
pad=args.padding,
remove=args.remove)
folds = []
best_losses = []
best_scores = []
for fold, ((train_img_paths, val_img_paths),
(train_labels, val_labels)) in enumerate(zip(img_paths,
labels)):
print('Fold [%d/%d]' % (fold + 1, len(img_paths)))
if os.path.exists('models/%s/model_%d.pth' % (args.name, fold + 1)):
log = pd.read_csv('models/%s/log_%d.csv' % (args.name, fold + 1))
best_loss, best_score = log.loc[log['val_loss'].values.argmin(),
['val_loss', 'val_score']].values
folds.append(str(fold + 1))
best_losses.append(best_loss)
best_scores.append(best_score)
continue
if args.remove_duplicate:
md5_df = pd.read_csv('inputs/strMd5.csv')
duplicate_img_paths = aptos2019_dir + '/' + md5_df[
(md5_df.strMd5_count > 1)
& (~md5_df.diagnosis.isnull())]['id_code'].values + '.png'
print(duplicate_img_paths)
for duplicate_img_path in duplicate_img_paths:
train_labels = train_labels[
train_img_paths != duplicate_img_path]
train_img_paths = train_img_paths[
train_img_paths != duplicate_img_path]
val_labels = val_labels[val_img_paths != duplicate_img_path]
val_img_paths = val_img_paths[
val_img_paths != duplicate_img_path]
# Train.
train_set = Dataset(train_img_paths,
train_labels,
transform=train_transform)
_, class_sample_counts = np.unique(train_labels, return_counts=True)
_, vclass_sample_counts = np.unique(val_labels, return_counts=True)
print('Train:', class_sample_counts)
print('Validation:', vclass_sample_counts)
if args.class_aware:
print('Class Aware(N:P):', args.N_ratio, ':', args.P_ratio)
print('Number of Train samples:', args.num_train_sample)
print('Number of Validation samples:', args.num_val_sample)
print('------------------------------------------')
weights = 1. / torch.tensor(class_sample_counts, dtype=torch.float)
weights = np.array([1, 1, 1, 1, 1])
samples_weights = weights[train_labels]
vsamples_weights = weights[val_labels]
tsampler = WeightedRandomSampler(weights=samples_weights,
num_samples=args.num_train_sample,
replacement=True)
vsampler = WeightedRandomSampler(weights=vsamples_weights,
num_samples=args.num_val_sample,
replacement=True)
train_loader = torch.utils.data.DataLoader(
train_set,
batch_size=args.batch_size,
shuffle=False if args.class_aware else True,
num_workers=args.num_workers,
sampler=tsampler if args.class_aware else None)
val_set = Dataset(val_img_paths, val_labels, transform=val_transform)
val_loader = torch.utils.data.DataLoader(
val_set,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.num_workers,
sampler=vsampler if args.class_aware else None)
# Create model.
model = get_model(model_name=args.arch,
num_outputs=num_outputs,
freeze_bn=args.freeze_bn,
dropout_p=args.dropout_p)
if args.print_model:
print(model)
model = model.cuda()
if args.pretrained_model is not None:
model.load_state_dict(
torch.load('models/%s/model_%d.pth' %
(args.pretrained_model, fold + 1)))
# Print the model.
if args.optimizer == 'Adam':
optimizer = optim.Adam(filter(lambda p: p.requires_grad,
model.parameters()),
lr=args.lr,
weight_decay=args.weight_decay)
elif args.optimizer == 'AdamW':
optimizer = optim.AdamW(filter(lambda p: p.requires_grad,
model.parameters()),
lr=args.lr,
weight_decay=args.weight_decay)
elif args.optimizer == 'RAdam':
optimizer = RAdam(filter(lambda p: p.requires_grad,
model.parameters()),
lr=args.lr,
weight_decay=args.weight_decay)
elif args.optimizer == 'SGD':
optimizer = optim.SGD(filter(lambda p: p.requires_grad,
model.parameters()),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay,
nesterov=args.nesterov)
elif args.optimizer == 'Ranger':
optimizer = Ranger(filter(lambda p: p.requires_grad,
model.parameters()),
lr=args.lr,
weight_decay=args.weight_decay)
elif args.optimizer == 'RangerLars':
optimizer = RangerLars(filter(lambda p: p.requires_grad,
model.parameters()),
lr=args.lr,
weight_decay=args.weight_decay)
elif args.optimizer == 'SAM':
base_optimizer = RangerLars
optimizer = SAM(model.parameters(),
base_optimizer,
lr=args.lr,
weight_decay=args.weight_decay)
# if args.scheduler == 'CosineAnnealingLR':
# scheduler = lr_scheduler.CosineAnnealingLR(
# optimizer, T_max=args.epochs, eta_min=args.min_lr)
if args.scheduler == 'CosineAnnealingLR':
scheduler = lr_scheduler.CosineAnnealingWarmRestarts(
optimizer, T_0=5, T_mult=6, eta_min=args.min_lr)
elif args.scheduler == 'ReduceLROnPlateau':
scheduler = lr_scheduler.ReduceLROnPlateau(optimizer,
factor=args.factor,
patience=args.patience,
verbose=1,
min_lr=args.min_lr)
log = pd.DataFrame(
index=[],
columns=['epoch', 'loss', 'score', 'val_loss', 'val_score'])
log = {
'epoch': [],
'loss': [],
'score': [],
'val_loss': [],
'val_score': [],
}
best_loss = float('inf')
best_score = 0
for epoch in range(args.epochs):
print('Epoch [%d/%d]' % (epoch + 1, args.epochs))
# Train for one epoch.
train_loss, train_score = train(args, train_loader, model,
criterion, optimizer, epoch)
# Evaluate on validation set.
val_loss, val_score = validate(args, val_loader, model, criterion)
if args.scheduler == 'CosineAnnealingLR':
scheduler.step()
elif args.scheduler == 'ReduceLROnPlateau':
scheduler.step(val_loss)
print(
'loss %.4f - score %.4f - val_loss %.4f - val_score %.4f - best_score %.4f'
% (train_loss, train_score, val_loss, val_score, best_score))
log['epoch'].append(epoch)
log['loss'].append(train_loss)
log['score'].append(train_score)
log['val_loss'].append(val_loss)
log['val_score'].append(val_score)
pd.DataFrame(log).to_csv('models/%s/log_%d.csv' %
(args.name, fold + 1),
index=False)
if val_score >= best_score:
if (val_score / train_score) <= 1.6:
torch.save(
model.state_dict(),
'models/%s/model_%d.pth' % (args.name, fold + 1))
best_loss = val_loss
best_score = val_score
print("=> saved best model")
print('val_loss: %f' % best_loss)
print('val_score: %f' % best_score)
folds.append(str(fold + 1))
best_losses.append(best_loss)
best_scores.append(best_score)
results = pd.DataFrame({
'fold':
folds + ['mean'],
'best_loss':
best_losses + [np.mean(best_losses)],
'best_score':
best_scores + [np.mean(best_scores)],
})
print(results)
results.to_csv('models/%s/results.csv' % args.name, index=False)
torch.cuda.empty_cache()
if not args.cv:
break
def get(conf, trial, param_gen):
SIZE = conf["input_size"]
# TRANSFORM
trans = [
transforms.Resize((SIZE, SIZE)),
transforms.CenterCrop(SIZE),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
]
trans = transforms.Compose(trans)
if conf["use_augmentation"]:
train_trans = [
transforms.Resize((SIZE, SIZE)),
transforms.CenterCrop(SIZE)
]
# Flip
if param_gen.suggest_categorical("augment_flip", [0, 1]) > 0:
train_trans.append(transforms.RandomHorizontalFlip())
# Rotation, Shift, Scale
rot_degree = param_gen.suggest_uniform("augment_rot", 0.0, 180.0)
translate = param_gen.suggest_uniform("augment_trans", 0.0, 0.3)
scale = param_gen.suggest_uniform("augment_scale", 1.0, 1.3)
train_trans.append(
transforms.RandomAffine(rot_degree,
translate=(translate, translate),
scale=(1 / scale, scale)))
train_trans += [
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
]
# Erasing
if param_gen.suggest_categorical("augment_erase", [0, 1]) > 0:
train_trans.append(transforms.RandomErasing())
train_trans = transforms.Compose(train_trans)
elif conf["use_randaug"]:
_N = param_gen.suggest_int("randaug_N", 0, 10)
_M = param_gen.suggest_int("randaug_M", 0, 30)
gray_p = param_gen.suggest_uniform("gray_p", 0.0, 1.0)
train_trans = [
transforms.Resize((SIZE, SIZE)),
transforms.CenterCrop(SIZE),
RandAugment(_N, _M),
transforms.RandomGrayscale(gray_p),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
]
train_trans = transforms.Compose(train_trans)
else:
train_trans = trans
trans_dict = {"train": train_trans, "valid": trans}
# DATASET GENERATOR
label_dirnames = glob(os.path.join(conf["data_dir"], "*"))
labels = []
for dirname in label_dirnames:
label = dirname.split(os.path.sep)[-1]
labels.append(label)
labels = sorted(labels)
def dataset_fold_generator(all=False):
folds = KFold(n_splits=conf["n_fold"], shuffle=True, random_state=1234)
dummy = np.arange(len(labels), dtype=np.int32)[:, np.newaxis]
for train_lab, dev_lab in list(
folds.split(dummy))[:conf["fold_trials"]]:
train_lab = set([labels[_] for _ in train_lab])
dev_lab = set([labels[_] for _ in dev_lab])
train_dataset = ImageFolder(root=conf["data_dir"],
transform=trans_dict["train"],
is_valid_file=lambda fn: fn.split(
os.path.sep)[-2] in train_lab)
train_sampler = samplers.MPerClassSampler(
labels=train_dataset.targets,
m=param_gen.suggest_categorical("m", conf["m_cands"]),
batch_size=conf["batch_size"],
length_before_new_iter=conf["data_per_epoch"])
dev_dataset = ImageFolder(
root=conf["data_dir"],
transform=trans_dict["valid"],
is_valid_file=lambda fn: fn.split(os.path.sep)[-2] in dev_lab)
yield train_dataset, dev_dataset, train_sampler, conf["batch_size"]
# MODEL, OPTIMIZER & LOSS
def model_generator():
if conf["trunk_model"] in [
"ResNet-18", "ResNet-34", "ResNet-50", "ResNet-101",
"ResNet-152"
]:
depth = re.match("ResNet-(?P<depth>[0-9]+)",
conf["trunk_model"]).group("depth")
trunk = eval(f"models.resnet{depth}(pretrained=True)")
trunk_output_size = trunk.fc.in_features
trunk.fc = nn.Identity()
elif conf["trunk_model"].startswith("timm:"):
_model_name = conf["trunk_model"][5:]
trunk = timm.create_model(_model_name, pretrained=True)
trunk.reset_classifier(0)
trunk_output_size = trunk.num_features
p_dropout = param_gen.suggest_uniform("p_dropout", 0.0, 1.0)
embedder = nn.Sequential(nn.Linear(trunk_output_size, conf["dim"]),
nn.Dropout(p_dropout))
classifier = nn.Linear(conf["dim"], len(labels), bias=False)
trunk.to("cuda")
embedder.to("cuda")
classifier.to("cuda")
model_dict = {
"trunk": trunk,
"embedder": embedder,
"classifier": classifier
}
#lr = param_gen.suggest_loguniform("model_lr", 1e-5, 1e-3)
# beta1 = 0.9
# beta2 = 0.999
# eps = 1e-8
decay = param_gen.suggest_loguniform("model_decay", 1e-10, 1e-2)
#lr = param_gen.suggest_loguniform("model_lr", 1e-5, 1e-2)
lr = param_gen.suggest_loguniform("model_lr", 10**(-4.75), 10**(-3.75))
beta1 = 1. - param_gen.suggest_loguniform("model_beta1", 1e-3, 1.)
beta2 = 1. - param_gen.suggest_loguniform("model_beta2", 1e-4, 1.)
eps = param_gen.suggest_loguniform("model_eps", 1e-8, 1)
optim_dict = {
"trunk_optimizer":
RAdam(trunk.parameters(),
lr=lr,
weight_decay=decay,
betas=(beta1, beta2),
eps=eps),
"embedder_optimizer":
RAdam(embedder.parameters(),
lr=lr,
weight_decay=decay,
betas=(beta1, beta2),
eps=eps),
"classifier_optimizer":
RAdam(classifier.parameters(),
lr=lr,
weight_decay=decay,
betas=(beta1, beta2),
eps=eps)
}
loss_type = conf["loss_type"]
metric_loss_info = ALL_LOSSES[loss_type](param_gen,
num_classes=len(labels),
embedding_size=conf["dim"])
if "param" in metric_loss_info:
# Add optimizer for loss
lr = param_gen.suggest_loguniform("loss_lr", 1e-5, 1e-2)
# beta1 = 1. - param_gen.suggest_loguniform("loss_beta1", 1e-3, 1.)
# beta2 = 1. - param_gen.suggest_loguniform("loss_beta2", 1e-4, 1.)
# eps = param_gen.suggest_loguniform("loss_eps", 1e-10, 1e-5)
beta1 = 0.9
beta2 = 0.999
eps = 1e-8
opt = RAdam(metric_loss_info["loss"].parameters(),
lr=lr,
betas=(beta1, beta2),
eps=eps)
optim_dict["metric_loss_optimizer"] = opt
loss_dict = {
"metric_loss": metric_loss_info["loss"],
"classifier_loss": torch.nn.CrossEntropyLoss()
}
#cls_weight = param_gen.suggest_loguniform("cls_weight", 0.01, 100.0)
cls_weight = param_gen.suggest_loguniform("cls_weight", 0.1, 1.0)
loss_weights = {"metric_loss": 1, "classifier_loss": cls_weight}
return {
"models": model_dict,
"optimizers": optim_dict,
"loss_funcs": loss_dict,
"loss_weights": loss_weights
}
return {
"modules": model_generator,
"fold_generator": dataset_fold_generator
}
def get_image_transform(mode, args, resize_size=256, crop_size=224):
if mode == 'weak':
if args.use_rrc_on_wa:
if args.custom_scale:
trfs = [
transforms.RandomResizedCrop(size=crop_size,
scale=(0.2, 1.0))
]
else:
trfs = [transforms.RandomResizedCrop(size=crop_size)]
else:
trfs = [
transforms.Resize((resize_size, resize_size)),
transforms.RandomCrop(crop_size),
]
trfs += [
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
]
elif mode == 'simclr':
s = args.aug_strength
color_jitter = transforms.ColorJitter(0.4 * s, 0.4 * s, 0.4 * s,
0.1 * s)
if args.use_rrc:
if args.custom_scale:
trfs = [
transforms.RandomResizedCrop(size=crop_size,
scale=(0.2, 1.0))
]
else:
trfs = [transforms.RandomResizedCrop(size=crop_size)]
else:
trfs = [
transforms.Resize((resize_size, resize_size)),
transforms.RandomCrop(crop_size),
]
trfs += [
transforms.RandomHorizontalFlip(),
]
if args.jitter:
trfs.append(transforms.RandomApply([color_jitter], p=0.8))
if args.grayscale:
trfs.append(transforms.RandomGrayscale(p=0.2))
if args.gaussblur:
trfs.append(GaussianBlur(kernel_size=int(0.1 * crop_size)))
trfs.append(transforms.ToTensor())
elif mode == 'randaug':
trfs = [RandAugment(args.ra_n, args.ra_m)]
if args.use_rrc:
if args.custom_scale:
trfs += [
transforms.RandomResizedCrop(size=crop_size,
scale=(0.2, 1.0))
]
else:
trfs += [transforms.RandomResizedCrop(size=crop_size)]
else:
trfs += [
transforms.Resize((resize_size, resize_size)),
transforms.RandomCrop(crop_size),
]
trfs += [
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
]
elif mode == 'test':
trfs = [
transforms.Resize((resize_size, resize_size)),
transforms.CenterCrop(crop_size),
transforms.ToTensor()
]
if args.norm_img:
if args.norm_img_mode == 'whitening':
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
elif args.norm_img_mode == 'pmone':
normalize = transforms.Normalize(mean=[0.5, 0.5, 0.5],
std=[0.5, 0.5, 0.5])
trfs.append(normalize)
return trfs
def get_dataset(dataset,
batch_size,
split=1.0,
cutout_length=0,
N=3,
M=5,
RandA=False):
"""
MEAN = [0.49139968, 0.48215827, 0.44653124]
STD = [0.24703233, 0.24348505, 0.26158768]
transf = [
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip()
]
normalize = [
transforms.Resize(224), # 32
transforms.ToTensor(),
transforms.Normalize(MEAN, STD)
]
cutout = []
if cutout_length > 0:
cutout.append(Cutout(cutout_length))
train_transform = transforms.Compose(transf + normalize + cutout)
# if RandA:
# Add RandAugment with N, M(hyperparameter)
# train_transform.transforms.insert(0, RandAugment(N, M))
valid_transform = transforms.Compose(normalize)
"""
train_transform = transforms.Compose([
# transforms.Resize((224, 224)), # 224 Spop
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465),
(0.2023, 0.1994, 0.2010)),
])
if RandA:
# Add RandAugment with N, M(hyperparameter)
train_transform.transforms.insert(0, RandAugment(N, M))
valid_transform = transforms.Compose([
# transforms.Resize((224, 224)), # 224 Spos
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465),
(0.2023, 0.1994, 0.2010)),
])
if dataset == "cifar10":
dataset_train = CIFAR10(root="./data",
train=True,
download=True,
transform=train_transform)
dataset_valid = CIFAR10(root="./data",
train=False,
download=True,
transform=valid_transform)
elif dataset == 'cifar100':
dataset_train = CIFAR100(root="./data",
train=True,
download=True,
transform=train_transform)
dataset_valid = CIFAR100(root="./data",
train=False,
download=True,
transform=valid_transform)
else:
raise NotImplementedError
split_idx = 0
train_sampler = None
if split < 1.0:
sss = StratifiedShuffleSplit(n_splits=5,
test_size=1 - split,
random_state=0)
sss = sss.split(list(range(len(dataset_train))), dataset_train.targets)
for _ in range(split_idx + 1):
train_idx, valid_idx = next(sss)
train_sampler = SubsetRandomSampler(train_idx)
valid_sampler = SubsetSampler(valid_idx)
else:
valid_sampler = SubsetSampler([])
train_loader = torch.utils.data.DataLoader(
dataset_train,
batch_size=batch_size,
shuffle=True if train_sampler is None else False,
num_workers=32,
pin_memory=True,
sampler=train_sampler,
drop_last=True)
# valid_loader = torch.utils.data.DataLoader(
# dataset_valid, batch_size=batch_size, shuffle=True if train_sampler is None else False, num_workers=16,
# pin_memory=True, sampler=valid_sampler, drop_last=True)
valid_loader = torch.utils.data.DataLoader(dataset_valid,
batch_size=batch_size,
shuffle=False,
num_workers=16,
pin_memory=True,
drop_last=True)
print('train/valid:{}/{},\t '
'batchsize:{} * train-step/valid-step:{}/{} -> train/valid:{}/{},\t'
'Split train/valid:{}/{} '.format(
len(dataset_train),
len(dataset_valid),
batch_size,
int(len(train_loader)),
int(len(valid_loader)),
int(len(train_loader)) * batch_size,
int(len(valid_loader)) * batch_size,
len(train_loader) / (len(dataset_train) // batch_size),
len(valid_loader) / (len(dataset_valid) // batch_size),
))
train_dataprovider = DataIterator(train_loader)
val_dataprovider = DataIterator(valid_loader)
return train_dataprovider, val_dataprovider, int(len(train_loader)), int(
len(valid_loader))
config = config(config_path)
exp_name = config['exp_name']
print(f"EXP: {exp_name}")
train_transforms = transforms.Compose([transforms.Resize((385, 505)),
transforms.RandomHorizontalFlip(),
transforms.RandomVerticalFlip(),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406],
[0.229, 0.224, 0.225])])
test_transforms = transforms.Compose([transforms.Resize((385, 505)),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406],
[0.229, 0.224, 0.225])])
train_transforms.transforms.insert(0, RandAugment(2, 14))
class2idx = config['class2idx']
batch = config['batch_size']
num_workers = config['num_workers']
train_data = datasets.ImageFolder(root = config['Train'], transform = train_transforms)
valid_data = myset(config['Val'], config['val_txt'], class2idx, test_transforms)
test_data = myset(config['Test'], config['test_txt'], class2idx, test_transforms)
train_loader = torch.utils.data.DataLoader(train_data, batch_size = batch , num_workers = num_workers,shuffle = True)
valid_loader = torch.utils.data.DataLoader(valid_data, batch_size = batch, num_workers = num_workers,shuffle = False)
test_loader = torch.utils.data.DataLoader(test_data, batch_size = batch, num_workers = num_workers,shuffle = False)
ms = []
model1 = res(3)
def __init__(self, img_size=(256, 256)):
self.randaugment = torchvision.transforms.Compose([
RandAugment(10, 10),
])
print("Augment ")
from torchvision.transforms import transforms
from RandAugment import RandAugment
import visdom
from PIL import Image
viz = visdom.Visdom()
transform_train = transforms.Compose([
transforms.RandomCrop(256),
# transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
])
# Add RandAugment with N, M(hyperparameter)
transform_train.transforms.insert(0, RandAugment(2, 10))
AB = Image.open("datasets/img2rain/trainA/norain-1.png").convert('RGB')
def cifar10_unsupervised_dataloaders():
print('Data Preparation')
train_transform = Compose([
Pad(4),
RandomCrop(32, fill=128),
RandomHorizontalFlip(),
ToTensor(),
Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
RandomErasing(scale=(0.1, 0.33)),
])
unsupervised_train_transformation = Compose([
Pad(4),
RandomCrop(32, fill=128),
ToTensor(),
Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
])
# RANDAUGMENT
unsupervised_train_transformation.transforms.insert(0, RandAugment(3, 9))
test_transform = Compose([
ToTensor(),
Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
])
# Train dataset with and without labels
cifar10_train_ds = datasets.CIFAR10('/data/', train=True, download=True)
num_classes = len(cifar10_train_ds.classes)
print('Loading dataset {0} for training -- Num_samples: {1}, num_classes: {2}'.format(datasets.CIFAR10.__name__,len(cifar10_train_ds),10))
labelled_indices = []
unlabelled_indices = []
indices = np.random.permutation(len(cifar10_train_ds))
class_counters = list([0] * num_classes)
max_counter = 10000 // num_classes
for i in indices:
dp = cifar10_train_ds[i]
if len(cifar10_train_ds) < sum(class_counters):
unlabelled_indices.append(i)
else:
y = dp[1]
c = class_counters[y]
if c < max_counter:
class_counters[y] += 1
labelled_indices.append(i)
else:
unlabelled_indices.append(i)
# Labelled and unlabelled dataset
train_labelled_ds = Subset(cifar10_train_ds, labelled_indices)
train_labelled_ds_t = AddTransform(train_labelled_ds, train_transform)
# unlabelled ds and aug ds
train_unlabelled_ds = Subset(cifar10_train_ds, unlabelled_indices)
train_unlabelled_ds = ConcatDataset([train_unlabelled_ds,train_labelled_ds])
# apply transformation for both
train_unlabelled_ds_t = AddTransform(train_unlabelled_ds, train_transform)
train_unlabelled_aug_ds_t = AddTransform(train_unlabelled_ds, unsupervised_train_transformation)
print('Labelled dataset -- Num_samples: {0}, classes: {1}, \n Unsupervised dataset -- Num_samples {2}, Augmentation -- Num_samples: {3}'
.format(len(train_labelled_ds_t), 10, len(train_unlabelled_ds_t), len(train_unlabelled_aug_ds_t)))
# Data loader for labeled and unlabeled train dataset
train_labelled = DataLoader(
train_labelled_ds_t,
batch_size=64,
shuffle=False,
num_workers=8,
pin_memory=True
)
train_unlabelled = DataLoader(
train_unlabelled_ds_t,
batch_size=64,
shuffle=False,
num_workers=8,
pin_memory=True
)
train_unlabelled_aug = DataLoader(
train_unlabelled_aug_ds_t,
batch_size=64,
shuffle=False,
num_workers=8,
pin_memory=True
)
# Data loader for test dataset
cifar10_test_ds = datasets.CIFAR10('/data/', transform=test_transform, train=False, download=True)
print('Test set -- Num_samples: {0}'.format(len(cifar10_test_ds)))
test = DataLoader(
cifar10_test_ds, batch_size=64,
shuffle=False,
num_workers=8,
pin_memory=True
)
return train_labelled, train_unlabelled, train_unlabelled_aug, test
def train():
data_transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.RandomHorizontalFlip(),
# ImageNetPolicy(),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
data_transform.transforms.insert(0, RandAugment(13, 3))
train_set = IMAGE_Dataset(Path(DATASET_ROOT_train), data_transform)
data_loader = DataLoader(dataset=train_set,
batch_size=16,
shuffle=True,
num_workers=1)
model = torch.hub.load('pytorch/vision:v0.6.0',
'googlenet',
pretrained=True)
# model = models.vgg19(pretrained=True)
# final_in_features = model.classifier[6].in_features
# model.classifier[6].out_features=train_set.num_classes
model = model.cuda(CUDA_DEVICES)
model.train()
best_model_params = copy.deepcopy(model.state_dict())
best_acc = 0.0
num_epochs = 200
criterion = nn.CrossEntropyLoss()
stepsize = 20
base_lr = 0.001
max_lr = 0.01
base_mm = 0.8
max_mm = 0.99
for epoch in range(num_epochs):
# newlr = get_triangular_lr(epoch,stepsize,base_lr,max_lr)
#mm=get_dynamic_momentum(epoch,stepsize,base_mm,max_mm)
optimizer = torch.optim.SGD(params=model.parameters(),
lr=0.001,
momentum=0.9)
print(f'Epoch: {epoch + 1}/{num_epochs}')
print('-' * len(f'Epoch: {epoch + 1}/{num_epochs}'))
training_loss = 0.0
training_corrects = 0
for i, (inputs, labels) in enumerate(data_loader):
inputs = Variable(inputs.cuda(CUDA_DEVICES))
labels = Variable(labels.cuda(CUDA_DEVICES))
optimizer.zero_grad()
outputs = model(inputs)
_, preds = torch.max(outputs.data, 1)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
training_loss += loss.item() * inputs.size(0)
# print(training_loss)
#revise loss.data[0]-->loss.item()
training_corrects += torch.sum(preds == labels.data)
#print(f'training_corrects: {training_corrects}')
training_loss = training_loss / len(train_set)
training_acc = training_corrects.double() / len(train_set)
print(
f'Training loss: {training_loss:.4f}\taccuracy: {training_acc:.4f}\n'
)
test_acc = test(model)
if test_acc > best_acc:
best_acc = test_acc
best_model_params = copy.deepcopy(model.state_dict())
model.load_state_dict(best_model_params)
torch.save(model, f'model-{best_acc:.02f}-best_train_acc.pth')
def __init__(self, mdlParams, indSet, index=None):
"""
Args:
mdlParams (dict): Configuration for loading
indSet (string): Indicates train, val, test
"""
# Mdlparams
self.mdlParams = mdlParams
# Current indSet = 'trainInd'/'valInd'/'testInd'
self.indSet = indSet
if self.indSet == 'trainInd':
self.root = mdlParams['dataDir'] + '/train'
self.index = index
elif self.indSet == 'valInd':
self.root = mdlParams['dataDir'] + '/train'
self.index = index
else:
self.root = mdlParams['dataDir'] + '/test'
self.names_list = []
# Number of classes
self.numClasses = mdlParams['numClasses']
# Size to crop
self.crop_size = mdlParams['crop_size']
# Model input size
self.input_size = (np.int32(mdlParams['input_size'][0]),
np.int32(mdlParams['input_size'][1]))
# Potential class balancing option
self.balancing = mdlParams['balance_classes']
# Potential setMean and setStd
self.setMean = mdlParams['setMean'].astype(np.float32)
self.setStd = mdlParams['setStd'].astype(np.float32)
# Only downsample
self.only_downsmaple = mdlParams.get('only_downsmaple', True)
# Meta csv
self.meta_path = mdlParams['meta_path']
self.meta_df = pd.read_pickle(self.meta_path)
self.subsets_size = [0, 0, 0, 0, 0, 0, 0]
self.image_path = []
self.image_labels = []
self.image_meta = []
for img in os.listdir(self.root):
id = img.split('.')[0]
label = list(self.meta_df.loc[self.meta_df['image_id'] == id,
'dx'])[0]
self.image_labels.append(CLASS_LABELS.index(label))
self.subsets_size[CLASS_LABELS.index(label)] += 1
self.image_meta.append(
self.meta_df.loc[self.meta_df['image_id'] == id,
self.mdlParams['meta_features']].to_numpy())
self.image_path.append(os.path.join(self.root, img))
self.image_path = np.asarray(self.image_path)
self.image_labels = np.asarray(self.image_labels)
self.image_meta = np.asarray(self.image_meta)
if indSet == 'trainInd':
if index is not None:
self.image_path = self.image_path[index]
self.image_labels = self.image_labels[index]
self.image_meta = self.image_meta[index]
all_transforms = []
if self.only_downsmaple:
all_transforms.append(transforms.Resize(self.input_size))
else:
all_transforms.append(transforms.Resize(self.crop_size))
all_transforms.append(transforms.RandomCrop(self.input_size))
if mdlParams.get('flip_lr_ud', False):
all_transforms.append(transforms.RandomHorizontalFlip())
all_transforms.append(transforms.RandomVerticalFlip())
# Full rot
if mdlParams.get('full_rot', 0) > 0:
if mdlParams.get('scale', False):
all_transforms.append(
transforms.RandomChoice([
transforms.RandomAffine(mdlParams['full_rot'],
scale=mdlParams['scale'],
shear=mdlParams.get(
'shear', 0),
resample=Image.NEAREST),
transforms.RandomAffine(mdlParams['full_rot'],
scale=mdlParams['scale'],
shear=mdlParams.get(
'shear', 0),
resample=Image.BICUBIC),
transforms.RandomAffine(mdlParams['full_rot'],
scale=mdlParams['scale'],
shear=mdlParams.get(
'shear', 0),
resample=Image.BILINEAR)
]))
else:
all_transforms.append(
transforms.RandomChoice([
transforms.RandomRotation(mdlParams['full_rot'],
resample=Image.NEAREST),
transforms.RandomRotation(mdlParams['full_rot'],
resample=Image.BICUBIC),
transforms.RandomRotation(mdlParams['full_rot'],
resample=Image.BILINEAR)
]))
# Color distortion
if mdlParams.get('full_color_distort') is not None:
all_transforms.append(
transforms.ColorJitter(
brightness=mdlParams.get('brightness_aug', 32. / 255.),
saturation=mdlParams.get('saturation_aug', 0.5),
contrast=mdlParams.get('contrast_aug', 0.5),
hue=mdlParams.get('hue_aug', 0.2)))
else:
all_transforms.append(
transforms.ColorJitter(brightness=32. / 255.,
saturation=0.5))
# Autoaugment
if self.mdlParams.get('randaugment', False):
all_transforms.append(
RandAugment(self.mdlParams.get('N'),
self.mdlParams.get('M')))
# Cutout
if self.mdlParams.get('cutout', 0) > 0:
all_transforms.append(
Cutout_v0(n_holes=1, length=self.mdlParams['cutout']))
# Normalize
all_transforms.append(transforms.ToTensor())
all_transforms.append(
transforms.Normalize(np.float32(self.mdlParams['setMean']),
np.float32(self.mdlParams['setStd'])))
# All transforms
self.train_composed = transforms.Compose(all_transforms)
self.valid_composed = transforms.Compose([
transforms.Resize(self.input_size),
transforms.ToTensor(),
transforms.Normalize(
torch.from_numpy(self.setMean).float(),
torch.from_numpy(self.setStd).float())
])
self.image_path = duplist(self.image_path)
self.image_labels = duplist(self.image_labels)
self.image_meta = duplist(self.image_meta)
else:
if index is not None:
self.image_path = self.image_path[index]
self.image_labels = self.image_labels[index]
self.image_meta = self.image_meta[index]
self.valid_composed = transforms.Compose([
transforms.Resize(self.input_size),
transforms.ToTensor(),
transforms.Normalize(
torch.from_numpy(self.setMean).float(),
torch.from_numpy(self.setStd).float())
])
def __init__(self, mdlParams, indSet, fold=-1):
"""
Args:
mdlParams (dict): Configuration for loading
indSet (string): Indicates train, val, test
"""
# Mdlparams
self.mdlParams = mdlParams
# Current indSet = 'trainInd'/'valInd'/'testInd'
self.indSet = indSet
if self.indSet == 'trainInd':
self.root = mdlParams['dataDir'] + '/train'
assert (fold > 0 and fold <= mdlParams['fold'])
self.fold = fold
elif self.indSet == 'valInd':
self.root = mdlParams['dataDir'] + '/train'
assert (fold > 0 and fold <= mdlParams['fold'])
self.fold = fold
else:
self.root = mdlParams['dataDir'] + '/test'
assert (fold < 0)
self.names_list = []
# Number of classes
self.numClasses = mdlParams['numClasses']
# Size to crop
self.crop_size = mdlParams['crop_size']
# Model input size
self.input_size = (np.int32(mdlParams['input_size'][0]),
np.int32(mdlParams['input_size'][1]))
# Potential class balancing option
self.balancing = mdlParams['balance_classes']
# Potential setMean and setStd
self.setMean = mdlParams['setMean'].astype(np.float32)
self.setStd = mdlParams['setStd'].astype(np.float32)
# Only downsample
self.only_downsmaple = mdlParams.get('only_downsmaple', True)
# Meta csv
self.meta_path = mdlParams['meta_path']
self.meta_df = pd.read_pickle(self.meta_path)
class_label = 0
self.subsets_size = []
self.image_path = []
for dir in os.listdir(self.root):
subset_img_path = []
subset_name_list = []
dir_path = os.path.join(self.root, dir)
for image in os.listdir(dir_path):
subset_name_list.append({
'id': image.split('.')[0],
'label': class_label,
})
subset_img_path.append(os.path.join(dir_path))
subset_size = len(subset_img_path)
self.names_list.append(subset_name_list)
self.image_path.append(subset_img_path)
self.subsets_size.append(subset_size)
class_label += 1
if indSet == 'trainInd':
new_names_list = []
new_image_path = []
for i in range(self.numClasses):
print('Before folding: ' + str(self.subsets_size[i]))
fold_len = self.subsets_size[i] // self.mdlParams['fold']
self.names_list[i] = self.names_list[i][:(
self.fold - 1) * fold_len] + self.names_list[i][self.fold *
fold_len:]
self.image_path[i] = self.image_path[i][:(
self.fold - 1) * fold_len] + self.image_path[i][self.fold *
fold_len:]
self.subsets_size[i] = len(self.names_list[i])
print('After folding: ' + str(self.subsets_size[i]))
new_names_list += self.names_list[i]
new_image_path += self.image_path[i]
self.names_list = new_names_list
self.image_path = new_image_path
all_transforms = []
if self.only_downsmaple:
all_transforms.append(transforms.Resize(self.input_size))
else:
all_transforms.append(transforms.Resize(self.crop_size))
all_transforms.append(transforms.RandomCrop(self.input_size))
if mdlParams.get('flip_lr_ud', False):
all_transforms.append(transforms.RandomHorizontalFlip())
all_transforms.append(transforms.RandomVerticalFlip())
# Full rot
if mdlParams.get('full_rot', 0) > 0:
if mdlParams.get('scale', False):
all_transforms.append(
transforms.RandomChoice([
transforms.RandomAffine(mdlParams['full_rot'],
scale=mdlParams['scale'],
shear=mdlParams.get(
'shear', 0),
resample=Image.NEAREST),
transforms.RandomAffine(mdlParams['full_rot'],
scale=mdlParams['scale'],
shear=mdlParams.get(
'shear', 0),
resample=Image.BICUBIC),
transforms.RandomAffine(mdlParams['full_rot'],
scale=mdlParams['scale'],
shear=mdlParams.get(
'shear', 0),
resample=Image.BILINEAR)
]))
else:
all_transforms.append(
transforms.RandomChoice([
transforms.RandomRotation(mdlParams['full_rot'],
resample=Image.NEAREST),
transforms.RandomRotation(mdlParams['full_rot'],
resample=Image.BICUBIC),
transforms.RandomRotation(mdlParams['full_rot'],
resample=Image.BILINEAR)
]))
# Color distortion
if mdlParams.get('full_color_distort') is not None:
all_transforms.append(
transforms.ColorJitter(
brightness=mdlParams.get('brightness_aug', 32. / 255.),
saturation=mdlParams.get('saturation_aug', 0.5),
contrast=mdlParams.get('contrast_aug', 0.5),
hue=mdlParams.get('hue_aug', 0.2)))
else:
all_transforms.append(
transforms.ColorJitter(brightness=32. / 255.,
saturation=0.5))
# Autoaugment
if self.mdlParams.get('randaugment', False):
all_transforms.append(
RandAugment(self.mdlParams.get('N'),
self.mdlParams.get('M')))
# Cutout
if self.mdlParams.get('cutout', 0) > 0:
all_transforms.append(
Cutout_v0(n_holes=1, length=self.mdlParams['cutout']))
# Normalize
all_transforms.append(transforms.ToTensor())
all_transforms.append(
transforms.Normalize(np.float32(self.mdlParams['setMean']),
np.float32(self.mdlParams['setStd'])))
# All transforms
self.composed = transforms.Compose(all_transforms)
elif indSet == 'validInd':
new_names_list = []
new_image_path = []
for i in range(self.numClasses):
print('Before folding: ' + str(self.subsets_size[i]))
fold_len = self.subsets_size[i] / self.fold
self.names_list[i] = self.names_list[i][(self.fold - 1) *
fold_len:self.fold *
fold_len]
self.image_path[i] = self.image_path[i][(self.fold - 1) *
fold_len:self.fold *
fold_len]
self.subsets_size[i] = len(self.names_list[i])
print('After folding: ' + str(self.subsets_size[i]))
new_names_list += self.names_list[i]
new_image_path += self.image_path[i]
self.names_list = new_names_list
self.image_path = new_image_path
self.composed = transforms.Compose([
transforms.Resize(self.input_size),
transforms.ToTensor(),
transforms.Normalize(
torch.from_numpy(self.setMean).float(),
torch.from_numpy(self.setStd).float())
])
else:
new_names_list = []
new_image_path = []
for i in range(self.numClasses):
new_names_list += self.names_list[i]
new_image_path += self.image_path[i]
self.names_list = new_names_list
self.image_path = new_image_path
self.composed = transforms.Compose([
transforms.Resize(self.input_size),
transforms.ToTensor(),
transforms.Normalize(
torch.from_numpy(self.setMean).float(),
torch.from_numpy(self.setStd).float())
])
from torchvision.transforms import transforms
from RandAugment import RandAugment
import cv2
import os
from PIL import Image
path = os.path.join(os.getcwd(), 'FMD', 'image', 'train')
transform_train = transforms.Compose([
transforms.RandomHorizontalFlip(),
# transforms.ToTensor(),
# transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
])
# Add RandAugment with N, M(hyperparameter)
transform_train.transforms.insert(0, RandAugment(1, 9))
for root, dirs, files in os.walk(path):
for f in files:
if f.endswith('.jpg'):
for i in range(10):
s = os.path.join(root, f)
img = cv2.imread(s)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img = Image.fromarray(img)
img = transform_train(img)
print(os.path.join(root, '{}_aug_{}'.format(i, f)))
img.save(os.path.join(root, '{}_aug_{}'.format(i, f)))
