def generator(
input_size=512,
batch_size=4,
train_list='/home/klara/klara/home/DeepSemanticText/resources/ims2.txt',
vis=False,
in_train=True,
geo_type=0):
image_list = np.array(get_images(train_list))
print('{} training images in {}'.format(image_list.shape[0], train_list))
index = np.arange(0, image_list.shape[0])
allow_empty = False
if not in_train:
allow_empty = True
transform = transforms.Compose([
transforms.ColorJitter(.3, .3, .3, .3),
transforms.RandomGrayscale(p=0.1)
])
while True:
if in_train:
np.random.shuffle(index)
images = []
image_fns = []
score_maps = []
geo_maps = []
training_masks = []
gtso = []
lbso = []
gt_idxs = []
im_id = 0
for i in index:
try:
im_name = image_list[i]
if in_train:
if random.uniform(0, 100) < 80:
im_name = image_list[int(
random.uniform(0,
min(19000, image_list.shape[0] -
1)))] #use with synthetic data
if not os.path.exists(im_name):
continue
im = cv2.imread(im_name)
if im is None:
continue
allow_empty = False
# print(im_name)
name = os.path.basename(im_name)
name = name[:-4]
# print im_fn
h, w, _ = im.shape
txt_fn = im_name.replace(
os.path.basename(im_name).split('.')[1], 'txt')
base_name = os.path.basename(txt_fn)
txt_fn_gt = '{0}/gt_{1}'.format(os.path.dirname(im_name),
base_name)
if (not (os.path.exists(txt_fn)
or os.path.exists(txt_fn_gt))) and not allow_empty:
continue
allow_empty = random.randint(0, 100) < 40
if os.path.exists(txt_fn_gt) and (
txt_fn_gt.find('/done/') != -1
or txt_fn_gt.find('/icdar-2015-Ch4/') != -1):
text_polys, text_tags, labels_txt = load_gt_annoataion(
txt_fn_gt,
txt_fn_gt.find('/icdar-2015-Ch4/') != -1)
elif os.path.exists(txt_fn) and (
txt_fn.find('/Latin/') != -1
or txt_fn.find('/Arabic/') != -1
or txt_fn.find('/Chinese/') != -1
or txt_fn.find('/Japanese/') != -1
or txt_fn.find('/Bangla/') != -1):
try:
text_polys, text_tags, labels_txt = load_annoataion(
txt_fn, im) # 读取文字的坐标信息和文本
except:
print(txt_fn)
import traceback
traceback.print_exc()
os.remove(im_name)
os.remove(txt_fn)
continue
else:
text_polys, text_tags, labels_txt = load_annoataion(
txt_fn, im)
if in_train:
if random.uniform(0, 100) < 50 or im.shape[
0] < 600 or im.shape[1] < 600: # 随机在周边填充
top = int(random.uniform(300, 500))
bottom = int(random.uniform(300, 500))
left = int(random.uniform(300, 500))
right = int(random.uniform(300, 500))
im = cv2.copyMakeBorder(im, top, bottom, left, right,
cv2.BORDER_CONSTANT)
if len(text_polys) > 0:
text_polys[:, :, 0] += left
text_polys[:, :, 1] += top
if random.uniform(0, 100) < 30 and False:
im = random_rotation(im, text_polys) # 随机旋转
if random.uniform(0, 100) < 30:
im = random_perspective(im, text_polys) # ？？？随机干哈
#im = random_crop(im, text_polys, vis=False)
scalex = random.uniform(0.5, 2) # 宽度和高度方向上随机比例
scaley = scalex * random.uniform(0.8, 1.2)
im = cv2.resize(im,
dsize=(int(im.shape[1] * scalex),
int(im.shape[0] * scaley)))
text_polys[:, :, 0] *= scalex
text_polys[:, :, 1] *= scaley
if random.randint(0, 100) < 10:
im = np.invert(im)
new_h, new_w, _ = im.shape
resize_h = input_size
resize_w = input_size
if input_size == -1:
image_size = [
im.shape[1] // 32 * 32, im.shape[0] // 32 * 32
]
while image_size[0] * image_size[1] > 1024 * 1024:
image_size[0] /= 1.2
image_size[1] /= 1.2
image_size[0] = int(image_size[0] // 32) * 32
image_size[1] = int(image_size[1] // 32) * 32
resize_h = int(image_size[1])
resize_w = int(image_size[0])
scaled = cut_image(im, (resize_w, resize_w),
text_polys) # 随机裁剪图片，裁剪后的尺寸为554,554
if scaled.shape[0] == 0 or scaled.shape[1] == 0:
continue
#transform_boxes(im, scaled, text_polys, text_tags, vis=False)
if scaled.shape[1] != resize_w or scaled.shape[0] != resize_h:
#continue
scalex = scaled.shape[1] / resize_w
scaley = scaled.shape[0] / resize_h
if scalex < 0.5 or scaley < 0.5:
continue
scaled = cv2.resize(scaled,
dsize=(int(resize_w), int(resize_h)))
if len(text_polys) > 0:
text_polys[:, :, 0] /= scalex
text_polys[:, :, 1] /= scaley
im = scaled
new_h, new_w, _ = im.shape
# pytorch的图像处理和变化
pim = PIL.Image.fromarray(np.uint8(im))
pim = transform(pim)
if use_pyblur == 1 and random.uniform(0, 100) < 30:
pim = RandomizedBlur(pim)
im = np.array(pim)
if geo_type == 0:
score_map, geo_map, training_mask, gt_idx, gt_out, labels_out = generate_rbox(
im, (new_h, new_w),
text_polys,
text_tags,
labels_txt,
vis=vis) # 获得分割的目标值（分类，上下左右，角度）
else:
score_map, geo_map, training_mask, gt_idx, gt_out, labels_out = generate_rbox2(
im, (new_h, new_w),
text_polys,
text_tags,
labels_txt,
vis=vis)
if score_map.sum() == 0 and (not allow_empty):
#print('empty image')
continue
image_fns.append(im_name)
images.append(im[:, :, :].astype(np.float32))
gtso.append(gt_out)
lbso.append(labels_out)
training_masks.append(training_mask)
score_maps.append(score_map)
gt_idxs.append(gt_idx)
geo_maps.append(geo_map)
im_id += 1
if len(images) == batch_size:
images = np.asarray(images, dtype=np.float)
images /= 128
images -= 1
training_masks = np.asarray(training_masks, dtype=np.uint8)
score_maps = np.asarray(score_maps, dtype=np.uint8)
geo_maps = np.asarray(geo_maps, dtype=np.float)
gt_idxs = np.asarray(gt_idxs, dtype=np.int)
yield images, image_fns, score_maps, geo_maps, training_masks, gtso, lbso, gt_idxs
images = []
image_fns = []
geo_maps = []
score_maps = []
geo_maps = []
training_masks = []
gtso = []
lbso = []
gt_idxs = []
im_id = 0
except Exception as e:
import traceback
traceback.print_exc()
continue
if not in_train:
print("finish")
yield None
break
log_interval = 10
val_interval = 1
# ============================ step 1/5 数据 ============================
split_dir = os.path.join("rmb_split")
train_dir = os.path.join(split_dir, "train")
valid_dir = os.path.join(split_dir, "valid")
norm_mean = [0.485, 0.456, 0.406]
norm_std = [0.229, 0.224, 0.225]
train_transform = transforms.Compose([
transforms.Resize((32, 32)),
transforms.RandomCrop(32, padding=4),
transforms.RandomGrayscale(p=0.9),
transforms.ToTensor(),
transforms.Normalize(norm_mean, norm_std),
])
valid_transform = transforms.Compose([
transforms.Resize((32, 32)),
transforms.ToTensor(),
transforms.Normalize(norm_mean, norm_std),
])
# 构建MyDataset实例
train_data = RMBDataset(data_dir=train_dir, transform=train_transform)
valid_data = RMBDataset(data_dir=valid_dir, transform=valid_transform)
# 构建DataLoder
def __init__(self,
mode,
image_size,
file_loc="data/train_new.csv",
shuffle=True,
fold=0,
do_transform=True):
self.mode = mode
self.image_size = image_size
self.do_transform = do_transform
df = pd.read_csv(file_loc)
cols_dict = dict(df.dtypes)
cols_dict = {k: str(v) for k, v in cols_dict.items()}
self.target_cols = [
x for x in cols_dict.keys()
if cols_dict[x] == 'int64' and x != 'fold'
]
self.img_folder = "train" if "fold" in df.columns else "test"
if self.mode == "train":
self.df = df[df["fold"] != fold]
else:
if self.img_folder == "train":
self.df = df[df["fold"] == fold]
else:
self.df = df
if shuffle:
self.df = self.df.sample(frac=1.0)
self.img_transform = transforms.Compose([
transforms.Resize((self.image_size, self.image_size)),
transforms.RandomGrayscale(p=1),
transforms.RandomHorizontalFlip(p=0.1),
transforms.RandomAffine(10,
translate=None,
scale=None,
shear=None,
resample=0,
fillcolor=0),
transforms.RandomPerspective(distortion_scale=0.2, p=0.1),
transforms.RandomRotation(10,
resample=False,
expand=False,
center=None,
fill=None),
transforms.ColorJitter(brightness=0.1,
contrast=0.1,
saturation=0.1,
hue=0.1),
transforms.ToTensor(),
transforms.RandomErasing(p=0.1,
scale=(0.02, 0.1),
ratio=(0.3, 3.3),
inplace=True),
])
self.weak_transform = transforms.Compose([
transforms.Resize((self.image_size, self.image_size)),
transforms.RandomGrayscale(p=1),
transforms.ToTensor()
])
def main_worker(gpu, ngpus_per_node, args):
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
print("=> creating model '{}'".format(args.arch))
model = moco.builder.MoCo(models.__dict__[args.arch], args.moco_dim,
args.moco_k, args.moco_m, args.moco_t, args.mlp)
print(model)
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):
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'])
message = f"=> loaded checkpoint '{args.resume}' " + \
f"(epoch {checkpoint['epoch']}, lr: {optimizer.param_groups[0]['lr']})"
print(message)
with open(args.logfile, 'a') as f:
print(message, file=f)
else:
print("=> no checkpoint found at '{}'".format(args.resume))
cudnn.benchmark = True
# Data loading code
traindir = os.path.join(args.data, 'base')
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
if args.aug_plus:
# MoCo v2's aug: similar to SimCLR https://arxiv.org/abs/2002.05709
augmentation = [
transforms.RandomResizedCrop(args.image_size, scale=(0.2, 1.)),
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([.1, 2.])],
p=0.5),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize
]
else:
# MoCo v1's aug: the same as InstDisc https://arxiv.org/abs/1805.01978
augmentation = [
transforms.RandomResizedCrop(args.image_size, scale=(0.2, 1.)),
transforms.RandomGrayscale(p=0.2),
transforms.ColorJitter(0.4, 0.4, 0.4, 0.4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(), normalize
]
train_dataset = datasets.ImageFolder(
traindir,
moco.loader.TwoCropsTransform(transforms.Compose(augmentation)))
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,
drop_last=True)
max_top1Acc = 0.0
for epoch in range(args.start_epoch, args.start_epoch + args.epochs):
if args.distributed:
train_sampler.set_epoch(epoch)
adjust_learning_rate(optimizer, epoch, args)
# train for one epoch
avgloss, accTop1, accTop5 = train(train_loader, model, criterion,
optimizer, epoch, args)
with open(args.logfile, 'a') as f:
print(f'{avgloss:.4f}\t{accTop1:.2f}\t{accTop5:.2f}', file=f)
# if not args.multiprocessing_distributed or (args.multiprocessing_distributed
# and args.rank % ngpus_per_node == 0):
if (epoch + 1) % args.freq_save == 0:
filename = os.path.join(
args.checkpoint_dir,
f'checkpoint_{epoch+1:04d}_Top1_{accTop1:.2f}.pth.tar')
save_checkpoint(
{
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict(),
},
is_best=False,
filename=filename)
if max_top1Acc < accTop1:
max_top1Acc = accTop1
filename = os.path.join(args.checkpoint_dir,
f'checkpoint_best_model.pth.tar')
save_checkpoint(
{
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict(),
},
is_best=False,
filename=filename)
with open(args.logfile, 'a') as f:
print(
f'End time : {datetime.now():%Y-%m-%d} {datetime.now():%H:%M:%S}',
file=f)
def generator(
batch_size=4,
train_list='/home/klara/klara/home/DeepSemanticText/resources/ims2.txt',
in_train=True,
rgb=False,
norm_height=32):
image_list = np.array(get_images(train_list))
print('{} training images in {}'.format(image_list.shape[0], train_list))
index = np.arange(0, image_list.shape[0])
transform = transforms.Compose([
transforms.ColorJitter(.3, .3, .3, .3),
transforms.RandomGrayscale(p=0.1)
])
batch_sizes = []
cb = batch_size
for i in range(0, len(buckets)):
batch_sizes.append(cb)
if i % 10 == 0 and cb > 2:
cb /= 2
max_samples = len(image_list) - 1
bucket_images = []
bucket_labels = []
bucket_label_len = []
for b in range(0, len(buckets)):
bucket_images.append([])
bucket_labels.append([])
bucket_label_len.append([])
while True:
if in_train:
np.random.shuffle(index)
for i in index:
try:
image_name = image_list[i]
src_del = " "
spl = image_name.split(" ")
if len(spl) == 1:
spl = image_name.split(",")
src_del = ","
image_name = spl[0].strip()
gt_txt = ''
if len(spl) > 1:
gt_txt = ""
delim = ""
for k in range(1, len(spl)):
gt_txt += delim + spl[k]
delim = src_del
if len(gt_txt
) > 1 and gt_txt[0] == '"' and gt_txt[-1] == '"':
gt_txt = gt_txt[1:-1]
if len(gt_txt) == 0:
continue
if image_name[len(image_name) - 1] == ',':
image_name = image_name[0:-1]
if not os.path.exists(image_name):
continue
if rgb:
im = cv2.imread(image_name)
else:
im = cv2.imread(image_name, cv2.IMREAD_GRAYSCALE)
if im is None:
continue
if image_name.find('/chinese_0/') != -1:
im = cv2.rotate(im, cv2.ROTATE_90_COUNTERCLOCKWISE
) # horizontal chinese text
if im.shape[0] > im.shape[1] and len(gt_txt) > 4:
# cv2.imshow('bad', im)
# print(image_name)
# cv2.waitKey(0)
continue
scale = norm_height / float(im.shape[0])
width = int(im.shape[1] * scale) + random.randint(
-2 * norm_height, 2 * norm_height)
best_diff = width
bestb = 0
for b in range(0, len(buckets)):
if best_diff > abs(width - buckets[b]):
best_diff = abs(width - buckets[b])
bestb = b
if random.randint(0, 100) < 10:
bestb += random.randint(-1, 1)
bestb = max(0, bestb)
bestb = min(bestb, (len(buckets) - 1))
width = buckets[bestb]
im = cv2.resize(im, (int(buckets[bestb]), norm_height))
if not rgb:
im = im.reshape(im.shape[0], im.shape[1], 1)
if in_train:
if random.randint(0, 100) < 10:
im = np.invert(im)
if not use_pyblur and random.randint(0, 100) < 10:
im = cv2.blur(im, (3, 3))
if not rgb:
im = im.reshape(im.shape[0], im.shape[1], 1)
if random.randint(0, 100) < 10:
warp_mat = cv2.getRotationMatrix2D(
(im.shape[1] / 2, im.shape[0] / 2), 0, 1)
warp_mat[0, 1] = random.uniform(-0.1, 0.1)
im = cv2.warpAffine(im, warp_mat,
(im.shape[1], im.shape[0]))
pim = PIL.Image.fromarray(np.uint8(im))
pim = transform(pim)
if use_pyblur:
if random.randint(0, 100) < 10:
pim = RandomizedBlur(pim)
im = np.array(pim)
bucket_images[bestb].append(im[:, :, :].astype(np.float32))
gt_labels = []
for k in range(len(gt_txt)):
if gt_txt[k] in codec_rev:
gt_labels.append(codec_rev[gt_txt[k]])
else:
print('Unknown char: {0}'.format(gt_txt[k]))
gt_labels.append(3)
if 'ARABIC' in ud.name(gt_txt[0]):
gt_labels = gt_labels[::-1]
bucket_labels[bestb].extend(gt_labels)
bucket_label_len[bestb].append(len(gt_labels))
if len(bucket_images[bestb]) == batch_sizes[bestb]:
images = np.asarray(bucket_images[bestb], dtype=np.float)
images /= 128
images -= 1
yield images, bucket_labels[bestb], bucket_label_len[bestb]
max_samples += 1
max_samples = min(max_samples, len(image_list) - 1)
bucket_images[bestb] = []
bucket_labels[bestb] = []
bucket_label_len[bestb] = []
except Exception as e:
import traceback
traceback.print_exc()
continue
if not in_train:
print("finish")
yield None
break
def main():
args = parse_option()
if args.gpu is not None:
print("Use GPU: {} for training".format(args.gpu))
# set the data loader
data_folder = os.path.join(args.data_folder, 'train')
# data_folder = args.data_folder
image_size = 224
mean = [0.485, 0.456, 0.406]
std = [0.229, 0.224, 0.225]
normalize = transforms.Normalize(mean=mean, std=std)
# Data augmentation
if args.aug == 'NULL':
train_transform = transforms.Compose([
transforms.RandomResizedCrop(image_size, scale=(args.crop, 1.)),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
])
elif args.aug == 'CJ':
train_transform = transforms.Compose([
transforms.RandomResizedCrop(image_size, scale=(args.crop, 1.)),
transforms.RandomGrayscale(p=0.2),
transforms.ColorJitter(0.4, 0.4, 0.4, 0.4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
])
else:
raise NotImplementedError('augmentation not supported: {}'.format(
args.aug))
# TOMO
# train_dataset = ImageFolderInstance(data_folder, transform=train_transform, two_crop=True)#args.moco)
train_dataset = config.dataset['cls'](
transform=[config.transforms_q, config.transforms_k],
**config.dataset['params'])
train_sampler = None
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=config.batch_size,
shuffle=(train_sampler is None),
num_workers=config.workers,
pin_memory=True,
sampler=train_sampler)
# create model and optimizer
n_data = len(train_dataset)
if args.model == 'resnet50':
model_q = config.model['cls'](**config.model['params'])
model_k = config.model['cls'](**config.model['params'])
# model_q = InsResNet50()
# model_k = InsResNet50()
# elif args.model == 'resnet50x2':
# model_q = InsResNet50(width=2)
# model_k = InsResNet50(width=2)
# elif args.model == 'resnet50x4':
# model_q = InsResNet50(width=4)
# model_k = InsResNet50(width=4)
else:
raise NotImplementedError('model not supported {}'.format(args.model))
# copy weights from `model_q' to `model_k'
moment_update(model_q, model_k, 0)
# set the contrast memory (queue of K keys) and criterion
contrast = MemoryMoCo(config.model['params']['params']['num_classes'],
n_data, args.nce_k, args.nce_t,
args.softmax).cuda(args.gpu)
criterion = NCESoftmaxLoss() if args.softmax else NCECriterion(n_data)
criterion = criterion.cuda(args.gpu)
model_q = model_q.cuda()
model_k = model_k.cuda()
optimizer = torch.optim.SGD(model_q.parameters(),
lr=args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay)
cudnn.benchmark = True
'''
# mixed precision training, speeds up training, however is likely to harm the downstream classification.
if args.amp:
model_q, optimizer = amp.initialize(model_q, optimizer, opt_level=args.opt_level)
optimizer_ema = torch.optim.SGD(model_k.parameters(),
lr=0,
momentum=0,
weight_decay=0)
model_k, optimizer_ema = amp.initialize(model_k, optimizer_ema, opt_level=args.opt_level)
'''
args.start_epoch = 1
'''
# optionally resume from a checkpoint
if args.resume:
if os.path.isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume, map_location='cpu')
# checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch'] + 1
model_q.load_state_dict(checkpoint['model_q'])
optimizer.load_state_dict(checkpoint['optimizer'])
contrast.load_state_dict(checkpoint['contrast'])
model_k.load_state_dict(checkpoint['model_k'])
if args.amp and checkpoint['opt'].amp:
print('==> resuming amp state_dict')
amp.load_state_dict(checkpoint['amp'])
print("=> loaded successfully '{}' (epoch {})"
.format(args.resume, checkpoint['epoch']))
del checkpoint
torch.cuda.empty_cache()
else:
print("=> no checkpoint found at '{}'".format(args.resume))
'''
# tensorboard
logger = tb_logger.Logger(logdir=args.tb_folder, flush_secs=2)
# routine
for epoch in range(args.start_epoch, args.epochs + 1):
adjust_learning_rate(epoch, args, optimizer)
print("==> training...")
time1 = time.time()
loss, prob = train_moco(epoch, train_loader, model_q, model_k,
contrast, criterion, optimizer, args)
time2 = time.time()
print('epoch {}, total time {:.2f}'.format(epoch, time2 - time1))
# tensorboard logger
logger.log_value('ins_loss', loss, epoch)
logger.log_value('ins_prob', prob, epoch)
logger.log_value('learning_rate', optimizer.param_groups[0]['lr'],
epoch)
# save model
if epoch % args.save_freq == 0:
print('==> Saving...')
state = {
'opt': args,
'model_q': model_q.state_dict(),
'contrast': contrast.state_dict(),
'optimizer': optimizer.state_dict(),
'epoch': epoch,
}
state['model_k'] = model_k.state_dict()
if args.amp:
state['amp'] = amp.state_dict()
save_file = os.path.join(
args.model_folder,
'ckpt_epoch_{epoch}.pth'.format(epoch=epoch))
torch.save(state, save_file)
# help release GPU memory
del state
torch.cuda.empty_cache()
def __init__(self, config):
super().__init__(config)
input_size = config.CONFIG_OPTIONS.get(
P.KEY_INPUT_SHAPE,
P.GLB_PARAMS[P.KEY_DATASET_METADATA][P.KEY_DS_INPUT_SHAPE])[1]
rel_delta = config.CONFIG_OPTIONS.get(P.KEY_DA_REL_DELTA, 0.25)
delta = int(rel_delta * input_size)
jit_brightness = config.CONFIG_OPTIONS.get(P.KEY_DA_JIT_BRIGHTNESS,
0.1)
jit_contrast = config.CONFIG_OPTIONS.get(P.KEY_DA_JIT_CONTRAST, 0.1)
jit_saturation = config.CONFIG_OPTIONS.get(P.KEY_DA_JIT_SATURATION,
0.1)
jit_hue = config.CONFIG_OPTIONS.get(P.KEY_DA_JIT_HUE, 20 / 360)
jit_p = config.CONFIG_OPTIONS.get(P.KEY_DA_JIT_P, 0.5)
grayscale_p = config.CONFIG_OPTIONS.get(P.KEY_DA_GREYSCALE_P, 0.2)
persp_scale = config.CONFIG_OPTIONS.get(P.KEY_DA_PERSP_SCALE, 0.25)
persp_p = config.CONFIG_OPTIONS.get(P.KEY_DA_PERSP_P, 0.3)
resize_p = config.CONFIG_OPTIONS.get(P.KEY_DA_RESIZE_P, 0.3)
rot_degrees = config.CONFIG_OPTIONS.get(P.KEY_DA_ROT_DEGREES, 180)
rot_p = config.CONFIG_OPTIONS.get(P.KEY_DA_ROT_P, 0.3)
transl_p = config.CONFIG_OPTIONS.get(P.KEY_DA_TRANSL_P, 0.5)
# A textual summary of the transformation
self.TRANSFORM_SUMMARY = "delta" + str(delta) + "jit_brightness" + str(jit_brightness) + "jit_contrast" + str(jit_contrast) + \
"jit_saturation" + str(jit_saturation) + "jit_hue" + str(jit_hue) + "jit_p" + str(jit_p) + "grayscale_p" + str(grayscale_p) + \
"persp_scale" + str(persp_scale) + "persp_p" + str(persp_p) + "resize_p" + str(resize_p) + \
"rot_degrees" + str(rot_degrees) + "rot_p" + str(rot_p) + "transl_p" + str(transl_p)
T_augm = []
# Random blur
# Random noise
T_augm.append(
transforms.RandomApply([
transforms.ColorJitter(brightness=jit_brightness,
contrast=jit_contrast,
saturation=jit_saturation,
hue=jit_hue)
],
p=jit_p))
T_augm.append(transforms.RandomGrayscale(p=grayscale_p))
T_augm.append(transforms.RandomHorizontalFlip())
T_augm.append(
transforms.RandomApply(
[
transforms.Resize(
input_size + delta // 2
), # Random perspective tends to shrink the image, so we enlarge it beforehand
transforms.RandomPerspective(distortion_scale=persp_scale,
p=1.0)
],
p=persp_p))
T_augm.append(
transforms.RandomApply(
[
transforms.Lambda(
RandomResize(input_size - delta,
input_size + delta)) # Random rescale
],
p=resize_p))
T_augm.append(
transforms.RandomApply(
[transforms.RandomRotation(degrees=rot_degrees, expand=True)],
p=rot_p))
# Random occlusion
T_augm.append(
transforms.RandomApply(
[
transforms.CenterCrop(
input_size + delta), # Take a fixed-size central crop
transforms.RandomCrop(
input_size
) # Take a smaller fixed-size crop at random position (random translation)
],
p=transl_p))
self.T_augm = transforms.Compose(T_augm)
# 超参数设置
best_acc = 0 # best test accuracy
start_epoch = 0 # start from epoch 0 or last checkpoint epoch
train_loss_everyepoch = 0
train_accuracy_everyepoch = 0
test_loss_everyepoch = 0
test_accuracy_everyepoch = 0
# Data
print('==> Preparing data..')
# torchvision.transforms是pytorch中的图像预处理包 一般用Compose把多个步骤整合到一起：
transform_train = transforms.Compose([ # 通过compose将各个变换串联起来
transforms.RandomCrop(32, padding=4), # 先四周填充0，再把图像随机裁剪成32*32
transforms.RandomHorizontalFlip(), # 以0.5的概率水平翻转给定的PIL图像
# transforms.RandomAffine(5.0), # python2.7 没有
transforms.RandomGrayscale(p=0.1), # 依概率 p 将图片转换为灰度图
transforms.ColorJitter(brightness=0.3,
contrast=0.3,
saturation=0.4,
hue=0.4),
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465),
(0.2023, 0.1994, 0.2010)), # R,G,B每层的归一化用到的均值和方差
])
transform_test = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
])
## if you cannot run the program because of "OUT OF MEMORY", you can decrease the batch_size properly.
def get_dataset(args):
### color augmentation ###
color_jitter = transforms.ColorJitter(0.8 * args.color_jitter_strength,
0.8 * args.color_jitter_strength,
0.8 * args.color_jitter_strength,
0.2 * args.color_jitter_strength)
rnd_color_jitter = transforms.RandomApply([color_jitter], p=0.8)
rnd_gray = transforms.RandomGrayscale(p=0.2)
learning_type = args.train_type
if args.dataset == 'cifar-10':
if learning_type == 'contrastive':
transform_train = transforms.Compose([
rnd_color_jitter,
rnd_gray,
transforms.RandomHorizontalFlip(),
transforms.RandomResizedCrop(32),
transforms.ToTensor(),
])
transform_test = transform_train
elif learning_type == 'linear_eval':
transform_train = transforms.Compose([
rnd_color_jitter,
rnd_gray,
transforms.RandomHorizontalFlip(),
transforms.RandomResizedCrop(32),
transforms.ToTensor(),
])
transform_test = transforms.Compose([
transforms.ToTensor(),
])
elif learning_type == 'test':
transform_train = transforms.Compose([
transforms.RandomHorizontalFlip(),
transforms.RandomResizedCrop(32),
transforms.ToTensor(),
])
transform_test = transforms.Compose([
transforms.ToTensor(),
])
else:
raise NotImplementedError('wrong learning type')
train_dst = CIFAR10(root='./Dataset',
train=True,
download=True,
transform=transform_train,
contrastive_learning=learning_type)
val_dst = CIFAR10(root='./Dataset',
train=False,
download=True,
transform=transform_test,
contrastive_learning=learning_type)
if learning_type == 'contrastive':
train_sampler = torch.utils.data.distributed.DistributedSampler(
train_dst,
num_replicas=args.ngpu,
rank=args.local_rank,
)
train_loader = torch.utils.data.DataLoader(
train_dst,
batch_size=args.batch_size,
num_workers=4,
pin_memory=False,
shuffle=(train_sampler is None),
sampler=train_sampler,
)
val_loader = torch.utils.data.DataLoader(
val_dst,
batch_size=100,
num_workers=4,
pin_memory=False,
shuffle=False,
)
return train_loader, train_dst, val_loader, val_dst, train_sampler
else:
train_loader = torch.utils.data.DataLoader(
train_dst,
batch_size=args.batch_size,
shuffle=True,
num_workers=4)
val_batch = 100
val_loader = torch.utils.data.DataLoader(val_dst,
batch_size=val_batch,
shuffle=False,
num_workers=4)
return train_loader, train_dst, val_loader, val_dst
if args.dataset == 'cifar-100':
if learning_type == 'contrastive':
transform_train = transforms.Compose([
rnd_color_jitter, rnd_gray,
transforms.RandomHorizontalFlip(),
transforms.RandomResizedCrop(32),
transforms.ToTensor()
])
transform_test = transform_train
elif learning_type == 'linear_eval':
transform_train = transforms.Compose([
rnd_color_jitter, rnd_gray,
transforms.RandomHorizontalFlip(),
transforms.RandomResizedCrop(32),
transforms.ToTensor()
])
transform_test = transforms.Compose([transforms.ToTensor()])
elif learning_type == 'test':
transform_train = transforms.Compose([
transforms.RandomCrop(32,
padding=4), # Different from cifar-10
transforms.RandomHorizontalFlip(),
transforms.ToTensor()
])
transform_test = transforms.Compose([transforms.ToTensor()])
else:
raise NotImplementedError('wrong learning type')
train_dst = CIFAR100(root='./Dataset',
train=True,
download=True,
transform=transform_train,
contrastive_learning=learning_type)
val_dst = CIFAR100(root='./Dataset',
train=False,
download=True,
transform=transform_test,
contrastive_learning=learning_type)
if learning_type == 'contrastive':
train_sampler = torch.utils.data.distributed.DistributedSampler(
train_dst,
num_replicas=args.ngpu,
rank=args.local_rank,
)
train_loader = torch.utils.data.DataLoader(
train_dst,
batch_size=args.batch_size,
num_workers=4,
pin_memory=True,
shuffle=(train_sampler is None),
sampler=train_sampler,
)
val_loader = torch.utils.data.DataLoader(
val_dst,
batch_size=100,
num_workers=4,
pin_memory=True,
)
return train_loader, train_dst, val_loader, val_dst, train_sampler
else:
train_loader = torch.utils.data.DataLoader(
train_dst,
batch_size=args.batch_size,
shuffle=True,
num_workers=4)
val_loader = torch.utils.data.DataLoader(val_dst,
batch_size=100,
shuffle=False,
num_workers=4)
return train_loader, train_dst, val_loader, val_dst
def parseTransform(self, names):
lst = []
for name in names.split('_'):
if "resize" in name:
value = name.strip("resize")
lst += [tt.Resize(self._parse_x(value))]
elif "normalize" in name:
value = self._parse_comma(name.strip("normalize"))
if len(value) == 2:
lst += [tt.Normalize(value[0:1], value[1:2])]
else:
lst += [tt.Normalize(value[0:3], value[3:6])]
elif "toTensor" in name:
lst += [tt.ToTensor()]
elif "randomCrop" in name:
value = name.strip("randomCrop")
lst += [tt.RandomCrop(self._parse_x(value))]
elif "centerCrop" in name:
value = name.strip("centerCrop")
lst += [tt.CenterCrop(self._parse_x(value))]
elif "randomVerticalFlip" in name:
lst += [tt.RandomVerticalFlip()]
elif "randomResizedCrop" in name:
v = self._parse_comma(name.strip("randomResizedCrop"))
lst += [
tt.RandomResizedCrop(int(v[0]), v[1:3], v[3:5], int(v[5]))
]
elif "grayScale" in name:
value = int(name.strip("grayScale"))
lst += [tt.grayScale(value)]
elif "randomRotation" in name:
striped_name = name.strip("randomRotation")
if ',' in striped_name:
v = self._parse_comma(striped_name)
else:
v = int(striped_name)
lst += [tt.RandomRotation(v)]
elif "randomGrayscale" in name:
v = float(value)
lst += [tt.RandomGrayscale(v)]
elif "toPILImage" in name:
lst += [tt.ToPILImage()]
elif "randomHorizontalFlip" in name:
lst += [tt.RandomHorizontalFlip()]
elif "randomVerticalFlip" in name:
lst += [tt.RandomVerticalFlip()]
elif "randomBackground" in name:
lst += [tt.RandomBackground(self.setting['data']['base'])]
elif "pad" in name:
striped_name = name.strip("pad")
v = int(striped_name)
lst += [tt.Pad(v)]
elif "randomPad" in name:
striped_name = name.strip("randomPad")
v = self._parse_comma(striped_name)
lst += [my_transforms.RandomPad(v[0], v[1])]
elif "colorJitter" in name:
v = self._parse_comma(name.strip("colorJitter"))
# default
# tt.ColorJitter(brightness=0, contrast=0, saturation=0, hue=0)
lst += [
tt.ColorJitter(brightness=v[0],
contrast=v[1],
saturation=v[2],
hue=v[3])
]
elif 'toNumpy' in name:
lst += [my_transforms.ToNumpy()]
elif 'correctExif' in name:
lst += [my_transforms.CorrectExif()]
elif 'randomAffine' in name:
value = self._parse_comma(name.strip('randomAffine'))
lst += [
tt.RandomAffine(degrees=value[0:2],
translate=value[2:4],
scale=value[4:6],
shear=value[6:8])
]
elif 'humanCrop' in name:
value = self._parse_comma(name.strip("humanCrop"))
lst += [
my_transforms.HumanCrop(margin=value[0],
weight_path=self.setting['data']
['base']['openPosePath'],
scale=value[1],
gpu_ids=str(int(value[2])))
]
elif 'toHSV' in name:
lst += [my_transforms.ToHSV()]
elif name == 'None':
pass
else:
raise NotImplementedError(
"{} is could not parse, this function is not implemented.".
format(name))
return lst
parser.add_argument('--exp', default='./cifar', type=str, help='experimentdir')
parser.add_argument('--type', default='10', type=int, help='cifar10 or 100')
args = parser.parse_args()
setup_runtime(2, [args.device])
device = 'cuda' if torch.cuda.is_available() else 'cpu'
knn_dim = 4096
best_acc = 0 # best test accuracy
start_epoch = 0 # start from epoch 0 or last checkpoint epoch
# Data
print('==> Preparing data..') #########################
transform_train = tfs.Compose([
tfs.Resize(256),
tfs.RandomResizedCrop(size=224, scale=(0.2, 1.)),
tfs.ColorJitter(0.4, 0.4, 0.4, 0.4),
tfs.RandomGrayscale(p=0.2),
tfs.RandomHorizontalFlip(),
tfs.ToTensor(),
tfs.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
])
transform_test = tfs.Compose([
tfs.Resize(256),
tfs.CenterCrop(224),
tfs.ToTensor(),
tfs.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
])
if args.type == 10:
trainset = CIFAR10Instance(root=args.datadir,
train=True,
download=True,
def main_worker(gpu, ngpus_per_node, args):
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
print("=> creating model '{}'".format(args.arch))
model = UniMoCo(models.__dict__[args.arch], args.moco_dim, args.moco_k,
args.moco_m, args.moco_t, args.mlp)
print(model)
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 = UnifiedContrastive().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'])
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(args.resume))
cudnn.benchmark = True
# Data loading code
traindir = os.path.join(args.data, 'train')
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
if args.aug_plus:
# MoCo v2's aug: similar to SimCLR https://arxiv.org/abs/2002.05709
augmentation = [
transforms.RandomResizedCrop(224, scale=(0.2, 1.)),
transforms.RandomApply(
[
transforms.ColorJitter(0.4, 0.4, 0.4,
0.1) # not strengthened
],
p=0.8),
transforms.RandomGrayscale(p=0.2),
transforms.RandomApply([GaussianBlur([.1, 2.])], p=0.5),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize
]
else:
# MoCo v1's aug: the same as InstDisc https://arxiv.org/abs/1805.01978
augmentation = [
transforms.RandomResizedCrop(224, scale=(0.2, 1.)),
transforms.RandomGrayscale(p=0.2),
transforms.ColorJitter(0.4, 0.4, 0.4, 0.4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(), normalize
]
train_dataset = datasets.ImageFolder(
traindir, TwoCropsTransform(transforms.Compose(augmentation)))
# build a semi-supervised dataset according to the file
if os.path.exists(args.supervised_list):
supervised_set = set()
for line in open(args.supervised_list, 'r'):
supervised_set.add(line.strip())
samples = []
for path, target in train_dataset.samples:
filename = os.path.basename(path)
if filename in supervised_set:
samples.append((path, target))
else:
samples.append((path, -1))
assert len(supervised_set) == len(
[label for _, label in samples if label != -1])
print(
f'Training with {len(supervised_set)} supervised images and {len(samples)-len(supervised_set)} unsupervised images.'
)
else:
print("Running with full superivsed reprensentation learning.")
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,
drop_last=True)
for epoch in range(args.start_epoch, args.epochs):
if args.distributed:
train_sampler.set_epoch(epoch)
adjust_learning_rate(optimizer, epoch, args)
# train for one epoch
train(train_loader, model, criterion, optimizer, epoch, args)
if not args.multiprocessing_distributed or (
args.multiprocessing_distributed
and args.rank % ngpus_per_node == 0):
save_checkpoint(
{
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict(),
},
is_best=False,
filename='checkpoint_{:04d}.pth.tar'.format(epoch))
def main():
global best_acc1
best_acc1 = 0
args = parse_option()
if args.gpu is not None:
print("Use GPU: {} for training".format(args.gpu))
# set the data loader
train_folder = os.path.join(args.data_folder, 'train')
val_folder = os.path.join(args.data_folder, 'val')
image_size = 224
crop_padding = 32
mean = [0.485, 0.456, 0.406]
std = [0.229, 0.224, 0.225]
normalize = transforms.Normalize(mean=mean, std=std)
if args.aug == 'NULL':
train_transform = transforms.Compose([
transforms.RandomResizedCrop(image_size, scale=(args.crop, 1.)),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
])
elif args.aug == 'CJ':
train_transform = transforms.Compose([
transforms.RandomResizedCrop(image_size, scale=(args.crop, 1.)),
transforms.RandomGrayscale(p=0.2),
transforms.ColorJitter(0.4, 0.4, 0.4, 0.4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
])
else:
raise NotImplementedError('augmentation not supported: {}'.format(
args.aug))
train_dataset = datasets.ImageFolder(train_folder, train_transform)
val_dataset = datasets.ImageFolder(
val_folder,
transforms.Compose([
transforms.Resize(image_size + crop_padding),
transforms.CenterCrop(image_size),
transforms.ToTensor(),
normalize,
]))
print(len(train_dataset))
train_sampler = None
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=(train_sampler is None),
num_workers=args.num_workers,
pin_memory=True,
sampler=train_sampler)
val_loader = torch.utils.data.DataLoader(val_dataset,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.num_workers,
pin_memory=True)
# create model and optimizer
if args.model == 'resnet50':
model = InsResNet50()
classifier = LinearClassifierResNet(args.layer, args.n_label, 'avg', 1)
elif args.model == 'resnet50x2':
model = InsResNet50(width=2)
classifier = LinearClassifierResNet(args.layer, args.n_label, 'avg', 2)
elif args.model == 'resnet50x4':
model = InsResNet50(width=4)
classifier = LinearClassifierResNet(args.layer, args.n_label, 'avg', 4)
else:
raise NotImplementedError('model not supported {}'.format(args.model))
print('==> loading pre-trained model')
ckpt = torch.load(args.model_path)
model.load_state_dict(ckpt['model'])
print("==> loaded checkpoint '{}' (epoch {})".format(
args.model_path, ckpt['epoch']))
print('==> done')
model = model.cuda()
classifier = classifier.cuda()
criterion = torch.nn.CrossEntropyLoss().cuda(args.gpu)
if not args.adam:
optimizer = torch.optim.SGD(classifier.parameters(),
lr=args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay)
else:
optimizer = torch.optim.Adam(classifier.parameters(),
lr=args.learning_rate,
betas=(args.beta1, args.beta2),
weight_decay=args.weight_decay,
eps=1e-8)
model.eval()
cudnn.benchmark = True
# set mixed precision training
# if args.amp:
# model = amp.initialize(model, opt_level=args.opt_level)
# classifier, optimizer = amp.initialize(classifier, optimizer, opt_level=args.opt_level)
# optionally resume from a checkpoint
args.start_epoch = 1
if args.resume:
if os.path.isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume, map_location='cpu')
# checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch'] + 1
classifier.load_state_dict(checkpoint['classifier'])
optimizer.load_state_dict(checkpoint['optimizer'])
best_acc1 = checkpoint['best_acc1']
best_acc1 = best_acc1.cuda()
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
if 'opt' in checkpoint.keys():
# resume optimization hyper-parameters
print('=> resume hyper parameters')
if 'bn' in vars(checkpoint['opt']):
print('using bn: ', checkpoint['opt'].bn)
if 'adam' in vars(checkpoint['opt']):
print('using adam: ', checkpoint['opt'].adam)
if 'cosine' in vars(checkpoint['opt']):
print('using cosine: ', checkpoint['opt'].cosine)
args.learning_rate = checkpoint['opt'].learning_rate
# args.lr_decay_epochs = checkpoint['opt'].lr_decay_epochs
args.lr_decay_rate = checkpoint['opt'].lr_decay_rate
args.momentum = checkpoint['opt'].momentum
args.weight_decay = checkpoint['opt'].weight_decay
args.beta1 = checkpoint['opt'].beta1
args.beta2 = checkpoint['opt'].beta2
del checkpoint
torch.cuda.empty_cache()
else:
print("=> no checkpoint found at '{}'".format(args.resume))
# set cosine annealing scheduler
if args.cosine:
# last_epoch = args.start_epoch - 2
# eta_min = args.learning_rate * (args.lr_decay_rate ** 3) * 0.1
# scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, args.epochs, eta_min, last_epoch)
eta_min = args.learning_rate * (args.lr_decay_rate**3) * 0.1
scheduler = optim.lr_scheduler.CosineAnnealingLR(
optimizer, args.epochs, eta_min, -1)
# dummy loop to catch up with current epoch
for i in range(1, args.start_epoch):
scheduler.step()
# tensorboard
logger = tb_logger.Logger(logdir=args.tb_folder, flush_secs=2)
# routine
for epoch in range(args.start_epoch, args.epochs + 1):
if args.cosine:
scheduler.step()
else:
adjust_learning_rate(epoch, args, optimizer)
print("==> training...")
time1 = time.time()
train_acc, train_acc5, train_loss = train(epoch, train_loader, model,
classifier, criterion,
optimizer, args)
time2 = time.time()
print('train epoch {}, total time {:.2f}'.format(epoch, time2 - time1))
logger.log_value('train_acc', train_acc, epoch)
logger.log_value('train_acc5', train_acc5, epoch)
logger.log_value('train_loss', train_loss, epoch)
logger.log_value('learning_rate', optimizer.param_groups[0]['lr'],
epoch)
print("==> testing...")
test_acc, test_acc5, test_loss = validate(val_loader, model,
classifier, criterion, args)
logger.log_value('test_acc', test_acc, epoch)
logger.log_value('test_acc5', test_acc5, epoch)
logger.log_value('test_loss', test_loss, epoch)
# save the best model
if test_acc > best_acc1:
best_acc1 = test_acc
state = {
'opt': args,
'epoch': epoch,
'classifier': classifier.state_dict(),
'best_acc1': best_acc1,
'optimizer': optimizer.state_dict(),
}
save_name = '{}_layer{}.pth'.format(args.model, args.layer)
save_name = os.path.join(args.save_folder, save_name)
print('saving best model!')
torch.save(state, save_name)
# save model
if epoch % args.save_freq == 0:
print('==> Saving...')
state = {
'opt': args,
'epoch': epoch,
'classifier': classifier.state_dict(),
'best_acc1': test_acc,
'optimizer': optimizer.state_dict(),
}
save_name = 'ckpt_epoch_{epoch}.pth'.format(epoch=epoch)
save_name = os.path.join(args.save_folder, save_name)
print('saving regular model!')
torch.save(state, save_name)
# tensorboard logger
pass
def main():
global args, sv_name, logs_dir, checkpoint_dir
write_arguments_to_file(
args, os.path.join('./', sv_name, sv_name + '_arguments.txt'))
use_cuda = torch.cuda.is_available()
if use_cuda:
torch.backends.cudnn.enabled = True
cudnn.benchmark = True
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
train_data_transform = transforms.Compose([
transforms.Resize((256, 256)),
transforms.RandomGrayscale(p=0.2),
transforms.ColorJitter(0.4, 0.4, 0.4, 0.4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(), normalize
])
val_data_transform = transforms.Compose(
[transforms.Resize((256, 256)),
transforms.ToTensor(), normalize])
train_dataGen = DataGeneratorSplitting(data=args.data,
dataset=args.dataset,
imgExt=args.imgEXT,
imgTransform=train_data_transform,
phase='train')
val_dataGen = DataGeneratorSplitting(data=args.data,
dataset=args.dataset,
imgExt=args.imgEXT,
imgTransform=val_data_transform,
phase='val')
train_data_loader = DataLoader(train_dataGen,
batch_size=args.batch_size,
num_workers=args.num_workers,
shuffle=True,
pin_memory=True)
val_data_loader = DataLoader(val_dataGen,
batch_size=args.batch_size,
num_workers=args.num_workers,
shuffle=False,
pin_memory=True)
trainloader_wo_shuf = DataLoader(train_dataGen,
batch_size=args.batch_size,
num_workers=args.num_workers,
shuffle=False,
pin_memory=True)
# create model and optimizer
n_data = len(train_dataGen)
model = ResNet18_cls(clsNum=len(train_dataGen.sceneList),
dim=args.dim).cuda()
model_ema = ResNet18_cls(clsNum=len(train_dataGen.sceneList),
dim=args.dim).cuda()
y_true = []
for idx, data in enumerate(
tqdm(trainloader_wo_shuf, desc="extracting training labels")):
label_batch = data['label'].to(torch.device("cpu"))
y_true += list(np.squeeze(label_batch.numpy()).astype(np.float32))
y_true = np.asarray(y_true)
CELoss = torch.nn.CrossEntropyLoss().cuda()
# copy weights from `model' to `model_ema'
moment_update(model, model_ema, 0)
criterion = NCACrossEntropy(torch.LongTensor(y_true),
args.margin / args.temperature).cuda()
optimizer = torch.optim.SGD(model.parameters(),
args.lr,
momentum=0.9,
weight_decay=1e-4,
nesterov=True)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=30, gamma=0.5)
train_writer = SummaryWriter(
os.path.join(logs_dir, 'runs', sv_name, 'training'))
val_writer = SummaryWriter(os.path.join(logs_dir, 'runs', sv_name, 'val'))
best_acc = 0
start_epoch = 0
if args.resume:
if os.path.isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
start_epoch = checkpoint['epoch'] + 1
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
lemniscate = checkpoint['lemniscate']
model_ema.load_state_dict(checkpoint['model_ema'])
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(args.resume))
else:
# define lemniscate and loss function (criterion)
ndata = len(train_dataGen)
lemniscate = LinearAverage(args.dim, ndata, args.temperature).cuda()
for epoch in range(start_epoch, args.epochs):
print('Epoch {}/{}'.format(epoch, args.epochs - 1))
print('-' * 10)
train_Moco(train_data_loader, model, model_ema, lemniscate, criterion,
CELoss, optimizer, epoch, train_writer)
acc = val(val_data_loader, trainloader_wo_shuf, model, epoch,
val_writer)
is_best_acc = acc > best_acc
best_acc = max(best_acc, acc)
save_checkpoint(
{
'epoch': epoch + 1,
# 'arch': args.arch,
'model_state_dict': model.state_dict(),
'lemniscate': lemniscate,
'best_acc': best_acc,
'optimizer': optimizer.state_dict(),
},
is_best_acc,
sv_name)
scheduler.step()
nn.Linear(100, 10))
model.to(device)
# Load Data
my_transforms = transforms.Compose(
[ # Compose makes it possible to have many transforms
transforms.Resize((36, 36)), # Resizes (32,32) to (36,36)
transforms.RandomCrop((32, 32)), # Takes a random (32,32) crop
transforms.ColorJitter(brightness=0.5), # Change brightness of image
transforms.RandomRotation(
degrees=45), # Perhaps a random rotation from -45 to 45 degrees
transforms.RandomHorizontalFlip(
p=0.5), # Flips the image horizontally with probability 0.5
transforms.RandomVerticalFlip(
p=0.05), # Flips image vertically with probability 0.05
transforms.RandomGrayscale(
p=0.2), # Converts to grayscale with probability 0.2
transforms.ToTensor(
), # Finally converts PIL image to tensor so we can train w. pytorch
transforms.Normalize(mean=[0.5, 0.5, 0.5],
std=[0.5, 0.5,
0.5]), # Note: these values aren't optimal
])
train_dataset = datasets.CIFAR10(root="dataset/",
train=True,
transform=my_transforms,
download=True)
train_loader = DataLoader(dataset=train_dataset,
batch_size=batch_size,
shuffle=True)
def main():
"""
Image Classification Network Trainer
"""
parser = train_args.get_args()
parser.add_argument('--version',
action='version',
version='%(prog)s ' + __version__ + ' by ' +
__author__)
cli_args = parser.parse_args()
# check for data directory
if not os.path.isdir(cli_args.data_directory):
print(f'Data directory {cli_args.data_directory} was not found.')
exit(1)
# check for save directory
if not os.path.isdir(cli_args.save_dir):
print(f'Directory {cli_args.save_dir} does not exist. Creating...')
os.makedirs(cli_args.save_dir)
# load categories
with open(cli_args.categories_json, 'r') as f:
cat_to_name = json.load(f)
# set output to the number of categories
output_size = len(cat_to_name)
print(f"Images are labeled with {output_size} categories.")
# prep data loader
expected_means = [0.485, 0.456, 0.406]
expected_std = [0.229, 0.224, 0.225]
max_image_size = 224
batch_size = 32
training_transforms = transforms.Compose([
transforms.RandomHorizontalFlip(p=0.25),
transforms.RandomRotation(25),
transforms.RandomGrayscale(p=0.02),
transforms.RandomResizedCrop(max_image_size),
transforms.ToTensor(),
transforms.Normalize(expected_means, expected_std)
])
training_dataset = datasets.ImageFolder(cli_args.data_directory,
transform=training_transforms)
training_dataloader = torch.utils.data.DataLoader(training_dataset,
batch_size=batch_size,
shuffle=True)
# Make model
if not cli_args.arch.startswith("vgg") and not cli_args.arch.startswith(
"densenet"):
print("Only supporting VGG and DenseNet")
exit(1)
print(f"Using a pre-trained {cli_args.arch} network.")
nn_model = models.__dict__[cli_args.arch](pretrained=True)
densenet_input = {
'densenet121': 1024,
'densenet169': 1664,
'densenet161': 2208,
'densenet201': 1920
}
input_size = 0
# Input size from current classifier if VGG
if cli_args.arch.startswith("vgg"):
input_size = nn_model.classifier[0].in_features
if cli_args.arch.startswith("densenet"):
input_size = densenet_input[cli_args.arch]
# Prevent back propagation on parameters
for param in nn_model.parameters():
param.requires_grad = False
od = OrderedDict()
hidden_sizes = cli_args.hidden_units
hidden_sizes.insert(0, input_size)
print(
f"Building a {len(cli_args.hidden_units)} hidden layer classifier with inputs {cli_args.hidden_units}"
)
for i in range(len(hidden_sizes) - 1):
od['fc' + str(i + 1)] = nn.Linear(hidden_sizes[i], hidden_sizes[i + 1])
od['relu' + str(i + 1)] = nn.ReLU()
od['dropout' + str(i + 1)] = nn.Dropout(p=0.15)
od['output'] = nn.Linear(hidden_sizes[i + 1], output_size)
od['softmax'] = nn.LogSoftmax(dim=1)
classifier = nn.Sequential(od)
# Replace classifier
nn_model.classifier = classifier
# Start clean by setting gradients of all parameters to zero.
nn_model.zero_grad()
# The negative log likelihood loss as criterion.
criterion = nn.NLLLoss()
# Adam: A Method for Stochastic Optimization
# https://arxiv.org/abs/1412.6980
print(f"Setting optimizer learning rate to {cli_args.learning_rate}.")
optimizer = optim.Adam(nn_model.classifier.parameters(),
lr=cli_args.learning_rate)
# Start with CPU
device = torch.device("cpu")
# Requested GPU
if cli_args.use_gpu and torch.cuda.is_available():
device = torch.device("cuda:0")
else:
print("GPU is not available. Using CPU.")
print(f"Sending model to device {device}.")
nn_model = nn_model.to(device)
data_set_len = len(training_dataloader.batch_sampler)
chk_every = 50
print(f'Using the {device} device to train.')
print(
f'Training on {data_set_len} batches of {training_dataloader.batch_size}.'
)
print(
f'Displaying average loss and accuracy for epoch every {chk_every} batches.'
)
for e in range(cli_args.epochs):
e_loss = 0
prev_chk = 0
total = 0
correct = 0
print(
f'\nEpoch {e+1} of {cli_args.epochs}\n----------------------------'
)
for ii, (images, labels) in enumerate(training_dataloader):
# Move images and labeles perferred device
# if they are not already there
images = images.to(device)
labels = labels.to(device)
# Set gradients of all parameters to zero.
optimizer.zero_grad()
# Propigate forward and backward
outputs = nn_model.forward(images)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
# Keep a running total of loss for
# this epoch
e_loss += loss.item()
# Accuracy
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
# Keep a running total of loss for
# this epoch
itr = (ii + 1)
if itr % chk_every == 0:
avg_loss = f'avg. loss: {e_loss/itr:.4f}'
acc = f'accuracy: {(correct/total) * 100:.2f}%'
print(
f' Batches {prev_chk:03} to {itr:03}: {avg_loss}, {acc}.')
prev_chk = (ii + 1)
print('Done... Saving')
nn_model.class_to_idx = training_dataset.class_to_idx
model_state = {
'epoch': cli_args.epochs,
'state_dict': nn_model.state_dict(),
'optimizer_dict': optimizer.state_dict(),
'classifier': nn_model.classifier,
'class_to_idx': nn_model.class_to_idx,
'arch': cli_args.arch
}
save_location = f'{cli_args.save_dir}/{cli_args.save_name}.pth'
print(f"Saving checkpoint to {save_location}")
torch.save(model_state, save_location)
def main():
parser = argparse.ArgumentParser(description='Image Classification.')
parser.add_argument('--model-name', type=str, default='resnet50')
parser.add_argument(
'--checkpoint-path',
type=str,
default='./tl/combine/sun_flow_combine_tl',
help=
'Path to save checkpoint, only the model with highest top1 acc will be saved,'
'And the records will also be writen in the folder')
parser.add_argument('--batch-size',
type=int,
default=128,
help='Batch size')
parser.add_argument('--lr',
type=float,
default=0.1,
help='Initial learning rate')
parser.add_argument('--epoch',
type=int,
default=100,
help='Maximum training epoch')
parser.add_argument('--start-epoch',
type=int,
default=0,
help='Start training epoch')
parser.add_argument('--root-dir',
type=str,
default='../data/',
help='path to the image folder')
parser.add_argument('--SUN-dir',
type=str,
default='./SUN397',
help='path to dataset')
parser.add_argument('--flower-dir',
type=str,
default='./tl-ssl/data/102Flower',
help='path to the image folder')
parser.add_argument(
'--train-json',
type=str,
default='../data/train2018.json',
help='path to the train folder, each class has a single folder')
parser.add_argument(
'--val-json',
type=str,
default='../data/val2018.json',
help='path to the validation folder, each class has a single folder')
# parser.add_argument('--test-dir', type=str, default='xxx/test',
# help='path to the train folder, each class has a single folder')
parser.add_argument('--cos',
type=bool,
default=False,
help='Use cos learning rate sheduler')
parser.add_argument('--schedule',
default=[50, 100, 150],
nargs='*',
type=int,
help='learning rate schedule (when to drop lr by 10x)')
parser.add_argument('--ratio', type=float, default=0.1)
parser.add_argument(
'--pretrained',
type=str,
default='Resume',
help='Load which pretrained model, '
'None : Do not load any weight, random initialize'
'Imagenet : official Imagenet pretrained model,'
'MoCo : Transfer model from Moco, path in $transfer-resume$'
'Transfer : Transfer model from Supervised pretrained, path in $transfer-resume$'
'Resume : Load checkpoint for corrupted training process, path in $resume$'
)
parser.add_argument('--transfer-resume',
type=str,
default='',
help='Path to load transfering pretrained model')
parser.add_argument('--resume',
type=str,
default='',
help='Path to resume a checkpoint')
parser.add_argument('--num-class',
type=int,
default=499,
help='Number of class for the classification')
parser.add_argument('--PRINT-INTERVAL',
type=int,
default=20,
help='Number of batch to print the loss')
args = parser.parse_args()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("Device {}".format(device))
# Create checkpoint file
if os.path.exists(args.checkpoint_path) == False:
os.makedirs(args.checkpoint_path)
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
test_trans = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop((224, 224)),
transforms.ToTensor(), normalize
])
augmentation = [
transforms.RandomResizedCrop(224, scale=(0.2, 1.)),
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([.1, 2.])], p=0.5),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize
]
flower = Flower102(root_dir=args.flower_dir, mode='train')
SUNdataset = SUNDataset(root=args.SUN_dir, number_per_class=1000)
trainset = CombineDataset(images1=flower.images,
images2=SUNdataset.images,
labels1=flower.labels,
labels2=SUNdataset.labels,
transforms=transforms.Compose(augmentation))
valset = Flower102(root_dir=args.flower_dir,
mode='val',
transforms=test_trans)
train_loader = DataLoader(trainset,
batch_size=args.batch_size,
num_workers=8,
shuffle=True,
pin_memory=True)
val_loader = DataLoader(valset,
batch_size=128,
num_workers=8,
pin_memory=True,
shuffle=True)
# test_loader = DataLoader(testset,batch_size=args.batch_size)
LOSS_FUNC = nn.CrossEntropyLoss().to(device)
print(args.model_name)
if args.pretrained == 'Imagenet':
# ImageNet supervised pretrained model
print('ImageNet supervised pretrained model')
model = MODEL_DICT[args.model_name](num_classes=args.num_class,
pretrained=True)
elif args.pretrained == 'MoCo':
# load weight from transfering model from moco
print('Load weight from transfering model from moco')
model = MODEL_DICT[args.model_name](num_classes=args.num_class,
pretrained=False)
if args.transfer_resume:
if os.path.isfile(args.transfer_resume):
print("=> loading checkpoint '{}'".format(
args.transfer_resume))
checkpoint = torch.load(args.pretrained, map_location="cpu")
# rename moco pre-trained keys
state_dict = checkpoint['state_dict']
for k in list(state_dict.keys()):
# retain only encoder_q up to before the embedding layer
if k.startswith('module.encoder_q') and not k.startswith(
'module.encoder_q.fc'):
# remove prefix
state_dict[
k[len("module.encoder_q."):]] = state_dict[k]
# delete renamed or unused k
del state_dict[k]
msg = model.load_state_dict(state_dict, strict=False)
assert set(msg.missing_keys) == {"fc.weight", "fc.bias"}
print("=> loaded pre-trained model '{}'".format(
args.transfer_resume))
else:
print("=> no checkpoint found at '{}'".format(
args.transfer_resume))
# freeze all layers but the last fc
for name, param in model.named_parameters():
if name not in ['fc.weight', 'fc.bias']:
param.requires_grad = False
# init the fc layer
model.fc.weight.data.normal_(mean=0.0, std=0.01)
model.fc.bias.data.zero_()
elif args.pretrained == 'Transfer':
# load weight from transfering model from supervised pretraining
print('Load weight from transfering model from supervised pretraining')
model = MODEL_DICT[args.model_name](num_classes=args.num_class,
pretrained=False)
if args.transfer_resume:
if os.path.isfile(args.transfer_resume):
print("=> loading checkpoint '{}'".format(
args.transfer_resume))
checkpoint = torch.load(args.transfer_resume)
msg = model.load_state_dict(checkpoint, strict=False)
assert set(msg.missing_keys) == {"fc.weight", "fc.bias"}
print("=> loaded checkpoint '{}' (epoch {})".format(
args.transfer_resume, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(
args.transfer_resume))
for name, param in model.named_parameters():
if name not in ['fc.weight', 'fc.bias']:
param.requires_grad = False
# init the fc layer
model.fc.weight.data.normal_(mean=0.0, std=0.01)
model.fc.bias.data.zero_()
else:
# Random Initialize
print('Random Initialize')
model = MODEL_DICT[args.model_name](num_classes=args.num_class,
pretrained=False)
# Dataparallel for multiple GPU usage
if torch.cuda.device_count() > 1:
print("Let's use", torch.cuda.device_count(), "GPUs!")
model = nn.DataParallel(model)
model = model.to(device)
# Optimizer and learning rate scheduler
optimizer = torch.optim.SGD(model.parameters(), lr=args.lr, momentum=0.9)
if args.pretrained == 'Resume':
# load weight from checkpoint
print('Load weight from checkpoint {}'.format(args.resume))
load_resume(args, model, optimizer, args.resume)
metric = []
for epoch in range(args.start_epoch, args.epoch):
adjust_learning_rate(optimizer, epoch, args)
train_loss = train(model, train_loader, optimizer, args.PRINT_INTERVAL,
epoch, args, LOSS_FUNC, device)
acc1, acc5, confusion_matrix, val_loss, aucs = test(
model, val_loader, args.num_class, LOSS_FUNC, device)
metric.append(acc1)
# Save train/val loss, acc1, acc5, confusion matrix(F1, recall, precision), AUCs
record = {
'epoch': epoch + 1,
'train loss': train_loss,
'val loss': val_loss,
'acc1': acc1,
'acc5': acc5,
'confusion matrix': confusion_matrix,
'AUCs': aucs
}
torch.save(
record,
os.path.join(args.checkpoint_path,
'recordEpoch{}.pth.tar'.format(epoch)))
# Only save the model with highest top1 acc
if np.max(metric) == acc1:
checkpoint = {
'epoch': epoch + 1,
'arch': args.model_name,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict(),
}
torch.save(checkpoint,
os.path.join(args.checkpoint_path, 'best.pth.tar'))
print("Model Saved")
def main():
global args, best_prec1
args = parser.parse_args()
args.distributed = args.world_size > 1
if args.distributed:
dist.init_process_group(backend=args.dist_backend,
init_method=args.dist_url,
world_size=args.world_size)
# create model
if args.pretrained:
print("=> using pre-trained model '{}'".format(args.arch))
model = models.__dict__[args.arch](pretrained=True)
else:
print("=> creating model '{}'".format(args.arch))
if not args.cifar:
model = models.__dict__[args.arch](low_dim=args.low_dim)
else:
model = models.resnet_cifar.__dict__[args.arch](
low_dim=args.low_dim)
if not args.distributed:
if args.arch.startswith('alexnet') or args.arch.startswith('vgg'):
model.features = torch.nn.DataParallel(model.features)
model.cuda()
else:
model = torch.nn.DataParallel(model).cuda()
else:
model.cuda()
model = torch.nn.parallel.DistributedDataParallel(model)
# Data loading code
'''
traindir = os.path.join(args.data, 'train')
valdir = os.path.join(args.data, 'val')
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
train_dataset = datasets.ImageFolderInstance(
traindir,
transforms.Compose([
transforms.RandomResizedCrop(224, scale=(0.2,1.)),
transforms.RandomGrayscale(p=0.2),
transforms.ColorJitter(0.4, 0.4, 0.4, 0.4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
]))
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)
val_loader = torch.utils.data.DataLoader(
datasets.ImageFolderInstance(valdir, transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
normalize,
])),
batch_size=args.batch_size, shuffle=False,
num_workers=args.workers, pin_memory=True)
'''
transform_train = transforms.Compose([
transforms.RandomResizedCrop(size=32, scale=(0.2, 1.)),
transforms.ColorJitter(0.4, 0.4, 0.4, 0.4),
transforms.RandomGrayscale(p=0.2),
#transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465),
(0.2023, 0.1994, 0.2010)),
])
transform_test = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465),
(0.2023, 0.1994, 0.2010)),
])
# normalize = transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
# transform = transforms.Compose([transforms.ToTensor(), normalize])
train_dataset = datasets.CIFAR10Instance(root=args.data,
train=True,
download=True,
transform=transform_train)
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=250,
shuffle=True,
num_workers=4,
pin_memory=True)
val_dataset = datasets.CIFAR10Instance(root=args.data,
train=False,
download=True,
transform=transform_test)
val_loader = torch.utils.data.DataLoader(val_dataset,
batch_size=250,
shuffle=False,
num_workers=4,
pin_memory=True)
classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse',
'ship', 'truck')
# define lemniscate and loss function (criterion)
ndata = train_dataset.__len__()
if args.nce_k > 0:
lemniscate = NCEAverage(args.low_dim, ndata, args.nce_k, args.nce_t,
args.nce_m)
criterion = NCECriterion(ndata).cuda()
else:
lemniscate = LinearAverage(args.low_dim, ndata, args.nce_t, args.nce_m)
criterion = nn.CrossEntropyLoss().cuda()
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))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
best_prec1 = checkpoint['best_prec1']
model.load_state_dict(checkpoint['state_dict'])
lemniscate = checkpoint['lemniscate']
optimizer.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(args.resume))
cudnn.benchmark = True
if args.evaluate:
kNN(0, model, lemniscate, train_loader, val_loader, 200, args.nce_t)
return
for epoch in range(args.start_epoch, args.epochs):
if args.distributed:
train_sampler.set_epoch(epoch)
adjust_learning_rate(optimizer, epoch)
if epoch % 10 == 0:
# remember best prec@1 and save checkpoint
prec1 = kNN(0, model, lemniscate, train_loader, val_loader, 200,
args.nce_t)
is_best = prec1 > best_prec1
best_prec1 = max(prec1, best_prec1)
save_checkpoint(
{
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': model.state_dict(),
'lemniscate': lemniscate,
'best_prec1': best_prec1,
'optimizer': optimizer.state_dict(),
}, is_best)
# evaluate KNN after last epoch
# train for one epoch
train(train_loader, model, lemniscate, criterion, optimizer, epoch)
def main():
global args, sv_name, logs_dir, checkpoint_dir
write_arguments_to_file(
args, os.path.join('./', sv_name, sv_name + '_arguments.txt'))
use_cuda = torch.cuda.is_available()
if use_cuda:
torch.backends.cudnn.enabled = True
cudnn.benchmark = True
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
train_data_transform = transforms.Compose([
transforms.Resize((256, 256)),
transforms.RandomGrayscale(p=0.2),
transforms.ColorJitter(0.4, 0.4, 0.4, 0.4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(), normalize
])
test_data_transform = transforms.Compose(
[transforms.Resize((256, 256)),
transforms.ToTensor(), normalize])
train_dataGen = DataGeneratorSM(data=args.data,
dataset=args.dataset,
train_per=args.train_per,
imgExt=args.imgEXT,
imgTransform=train_data_transform,
phase='train')
train_test_dataGen = DataGeneratorSM(data=args.data,
dataset=args.dataset,
train_per=args.train_per,
imgExt=args.imgEXT,
imgTransform=train_data_transform,
phase='test')
train_dataGen_ = DataGeneratorSM(data=args.data,
dataset=args.dataset,
train_per=args.train_per,
imgExt=args.imgEXT,
imgTransform=test_data_transform,
phase='train')
test_dataGen = DataGeneratorSM(data=args.data,
dataset=args.dataset,
train_per=args.train_per,
imgExt=args.imgEXT,
imgTransform=test_data_transform,
phase='test')
train_data_loader = DataLoader(train_dataGen,
batch_size=args.batch_size,
num_workers=args.num_workers,
shuffle=True,
pin_memory=True)
train_test_data_loader = DataLoader(train_test_dataGen,
batch_size=args.batch_size,
num_workers=args.num_workers,
shuffle=True,
pin_memory=True)
test_data_loader = DataLoader(test_dataGen,
batch_size=args.batch_size,
num_workers=args.num_workers,
shuffle=False,
pin_memory=True)
trainloader_wo_shuf = DataLoader(train_dataGen_,
batch_size=args.batch_size,
num_workers=args.num_workers,
shuffle=False,
pin_memory=True)
if args.MLP:
model = ResNet18_MLP(dim=args.dim).cuda()
else:
model = ResNet18(dim=args.dim).cuda()
loss_fn = NormSoftmaxLoss_Margin(args.dim,
len(train_dataGen.sceneList),
margin=args.margin,
temperature=args.temperature).cuda()
optimizer = torch.optim.SGD(list(model.parameters()) +
list(loss_fn.parameters()),
args.lr,
momentum=0.9,
weight_decay=1e-4,
nesterov=True)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=30, gamma=0.5)
best_acc = 0
start_epoch = 0
if args.resume:
if os.path.isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
start_epoch = checkpoint['epoch']
best_acc = checkpoint['best_acc']
model.load_state_dict(checkpoint['model_state_dict'])
loss_fn.load_state_dict(checkpoint['loss_state_dict'])
# lemniscate = checkpoint['lemniscate']
optimizer.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(args.resume))
train_writer = SummaryWriter(
os.path.join(logs_dir, 'runs', sv_name, 'training'))
val_writer = SummaryWriter(os.path.join(logs_dir, 'runs', sv_name, 'val'))
for epoch in range(start_epoch, args.epochs):
print('Epoch {}/{}'.format(epoch, args.epochs - 1))
print('-' * 10)
train(train_data_loader, train_test_data_loader, model, optimizer,
loss_fn, epoch, train_writer)
acc = val(test_data_loader, trainloader_wo_shuf, model, epoch,
val_writer)
is_best_acc = acc > best_acc
best_acc = max(best_acc, acc)
save_checkpoint(
{
'epoch': epoch + 1,
# 'arch': args.arch,
'model_state_dict': model.state_dict(),
'loss_state_dict': loss_fn.state_dict(),
# 'lemniscate': lemniscate,
'best_acc': best_acc,
'optimizer': optimizer.state_dict(),
},
is_best_acc,
sv_name)
scheduler.step()
def main_worker(gpu, ngpus_per_node, args):
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
model = sim2sem.builder.Sim2Sem(ResConvEncoder,
args.fea_dim,
args.num_clusters,
args.k,
args.m,
args.t,
args.sr,
freeze_encoder_s=True)
print(model)
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)
optimizer_pre = torch.optim.SGD(model.parameters(),
args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
# optionally resume from a checkpoint
args.resume = '{}/checkpoint_last.pth.tar'.format(args.save_folder)
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_pre = checkpoint['epoch_pre']
model.load_state_dict(checkpoint['state_dict'])
optimizer_pre.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch_pre']))
else:
print("=> no checkpoint found at '{}'".format(args.resume))
cudnn.benchmark = True
# Data loading code
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
if args.aug_plus:
# MoCo v2's aug: similar to SimCLR https://arxiv.org/abs/2002.05709
augmentation = [
transforms.RandomResizedCrop(224, scale=(0.2, 1.)),
transforms.RandomApply(
[
transforms.ColorJitter(0.4, 0.4, 0.4,
0.1) # not strengthened
],
p=0.8),
transforms.RandomGrayscale(p=0.2),
transforms.RandomApply([sim2sem.loader.GaussianBlur([.1, 2.])],
p=0.5),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize
]
else:
# MoCo v1's aug: the same as InstDisc https://arxiv.org/abs/1805.01978
augmentation = [
transforms.RandomResizedCrop(224, scale=(0.2, 1.)),
transforms.RandomGrayscale(p=0.2),
transforms.ColorJitter(0.4, 0.4, 0.4, 0.4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(), normalize
]
train_sim = STL10(args.data_folder,
split="train+test+unlabeled",
transform=sim2sem.loader.TwoCropsTransform(
transforms.Compose(augmentation)))
if args.distributed:
train_sampler_sim = torch.utils.data.distributed.DistributedSampler(
train_sim)
else:
train_sampler_sim = None
train_loader_sim = torch.utils.data.DataLoader(
train_sim,
batch_size=args.prebatch_size,
shuffle=(train_sampler_sim is None),
num_workers=args.workers,
pin_memory=True,
sampler=train_sampler_sim,
drop_last=True)
# Pretraining
for epoch_pre in range(args.start_epoch_pre, args.epochs_pre):
if args.distributed:
train_sampler_sim.set_epoch(epoch_pre)
adjust_learning_rate_pre(optimizer_pre, epoch_pre, args)
# train for one epoch
train_pre(train_loader_sim, model, criterion, optimizer_pre, epoch_pre,
args)
if not args.multiprocessing_distributed or (
args.multiprocessing_distributed
and args.rank % ngpus_per_node == 0 and
(epoch_pre + 1) % args.save_freq == 0):
save_checkpoint(
{
'epoch_pre': epoch_pre + 1,
'state_dict': model.state_dict(),
'optimizer': optimizer_pre.state_dict(),
},
is_best=False,
filename='{}/checkpoint_pre_{:04d}.pth.tar'.format(
args.save_folder, epoch_pre))
save_checkpoint(
{
'epoch_pre': epoch_pre + 1,
'state_dict': model.state_dict(),
'optimizer': optimizer_pre.state_dict(),
},
is_best=False,
filename='{}/checkpoint_last.pth.tar'.format(args.save_folder))
if (epoch_pre + 1) == args.epochs_pre:
save_checkpoint(
{
'epoch_pre': epoch_pre + 1,
'state_dict': model.state_dict(),
'optimizer': optimizer_pre.state_dict(),
},
is_best=False,
filename='{}/checkpoint_final.pth.tar'.format(
args.save_folder))
parser.add_argument('--cifar100', action='store_true')
parser.add_argument('--cifar10', action='store_true')
parser.add_argument('--miniimagenet', action='store_true')
parser.add_argument('--half', action='store_true')
parser.add_argument('--fine_labels', action='store_true') # only valid when using cifar100
parser.add_argument('--random_kmeans', action='store_true')
args = parser.parse_args()
n_tasks = args.num_tasks
# miniimagenet
transform_mini_train = transforms.Compose(
[
transforms.Resize(64),
transforms.RandomHorizontalFlip(),
transforms.RandomGrayscale(),
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))])
transform_mini_test = transforms.Compose(
[
transforms.Resize(64),
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))])
# cifar100
transform_cifar100_train = transforms.Compose(
[
transforms.RandomHorizontalFlip(),
transforms.RandomGrayscale(),
transforms.ToTensor(),
def main():
parser = argparse.ArgumentParser(
description='Scene Classification Training')
parser.add_argument('--device',
default='cuda:0',
type=str,
required=False,
help='GPU ids')
parser.add_argument('--epoch',
default=350,
type=int,
required=True,
help='training epochs')
parser.add_argument('--alldata',
dest='alldata',
action='store_true',
help='use alldata to train')
parser.add_argument('--lr', default=0.01, type=float, help='learning rate')
parser.add_argument('--weight_decay',
'-w',
default=5e-4,
type=float,
help='weight_decay')
parser.add_argument('--batch_size',
default=128,
type=int,
help='training batch size')
parser.add_argument('--output_dir', default='./checkpoint', type=str)
parser.add_argument('--warm_up_epochs', default=10, type=int)
parser.add_argument('--log_file', type=str, default='./log/default.log')
parser.add_argument('--params', type=str, default=None)
args = parser.parse_args()
log_file_name = args.log_file
logger.add(log_file_name)
logger.info('args:\n' + args.__repr__())
batch_size = args.batch_size
output_dir = args.output_dir
if not os.path.exists(output_dir):
os.makedirs(output_dir)
device = args.device
best_acc = 0 # best test accuracy
# Data
logger.info('==> Preparing data..')
transform_train = transforms.Compose([
transforms.ToPILImage(),
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ColorJitter(brightness=0.5,
contrast=0.5,
saturation=0.5,
hue=0.5),
transforms.RandomGrayscale(p=0.1),
transforms.ToTensor(),
])
transform_test = transforms.Compose([
transforms.ToPILImage(), #适应自己编写的dataset
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor()
])
train_dir = os.path.join(DATA_PATH, 'train/train')
if args.alldata == False:
whole_set = SceneDataset(
annotations_csv='/home/linhw/myproject/data_scene/train_labels.csv',
root_dir=train_dir,
transform=transform_train)
whole_set2 = SceneDataset(
annotations_csv='/home/linhw/myproject/data_scene/train_labels.csv',
root_dir=train_dir,
transform=transform_test)
whole_len = len(whole_set)
train_len = int(whole_len * 0.9) #划分train和val数据集
val_len = whole_len - train_len
indices = random.sample(range(0, whole_len), train_len)
indices2 = []
for x in range(0, whole_len):
if x not in indices:
indices2.append(x)
trainset = torch.utils.data.Subset(whole_set, indices)
testset = torch.utils.data.Subset(whole_set2, indices2)
else:
trainset = SceneDataset(
annotations_csv='/home/linhw/myproject/data_scene/train_labels.csv',
root_dir=train_dir,
transform=transform_train)
testset = torch.utils.data.Subset(
SceneDataset(annotations_csv=
'/home/linhw/myproject/data_scene/train_labels.csv',
root_dir=train_dir,
transform=transform_test), range(8000)
) # when whole data is True, randomly choose 8000 training samples. Test acc here is training acc.
trainloader = torch.utils.data.DataLoader(trainset,
batch_size=batch_size,
shuffle=True,
num_workers=8)
testloader = torch.utils.data.DataLoader(testset,
batch_size=batch_size,
shuffle=False,
num_workers=8)
# Model
logger.info('==> Building model..')
model = EfficientNet.from_name('efficientnet-b0', num_classes=100)
model = model.to(device)
if args.params is not None:
model.load_state_dict(torch.load(args.params, map_location=device))
optimizer = optim.SGD(model.parameters(),
lr=args.lr,
momentum=0.9,
weight_decay=args.weight_decay)
#optimizer = optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
warm_up_with_cosine_lr = lambda epoch: epoch / args.warm_up_epochs if epoch <= args.warm_up_epochs else 0.5 * (
math.cos((epoch - args.warm_up_epochs) /
(args.epoch - args.warm_up_epochs) * math.pi) + 1)
scheduler = torch.optim.lr_scheduler.LambdaLR(
optimizer, lr_lambda=warm_up_with_cosine_lr)
best_acc = 0
for epoch in range(args.epoch):
logger.info("Epoch {} started".format(epoch))
train_acc, training_loss = train(model, optimizer, trainloader, device)
logger.info(
"train acc = {:.4f}, training loss = {:.4f} lr = {:.4f}".format(
train_acc, training_loss, warm_up_with_cosine_lr(epoch)))
test_acc, test_loss = test(model, testloader, device)
logger.info("test acc = {:.4f}, test loss = {:.4f}".format(
test_acc, test_loss))
if test_acc > best_acc:
best_acc = test_acc
logger.info("best acc improved to {:.4f}".format(best_acc))
model_to_save = model.module if hasattr(model, 'module') else model
torch.save(model_to_save.state_dict(),
'{}/bset_model.pt'.format(output_dir))
logger.info("model saved to {}/bset_model.pt".format(output_dir))
model_to_save = model.module if hasattr(model, 'module') else model
torch.save(model_to_save.state_dict(),
'{}/last_model.pt'.format(output_dir))
logger.info("model saved to {}/last_model.pt".format(output_dir))
model_to_save = model.module if hasattr(model, 'module') else model
torch.save(model_to_save.state_dict(),
'{}/model.pt'.format(output_dir))
logger.info("model saved to {}/model.pt".format(output_dir))
scheduler.step()
logger.info("Epoch {} ended, best acc = {:.4f}".format(
epoch, best_acc))
logger.info("Training finished, best_acc = {:.4f}".format(best_acc))
def main(env, policy_net, target_net, action_space):
Transition = namedtuple(
"Transition", ('state', 'action', 'reward', 'next_state', "done"))
rm = experience_replay()
opt = torch.optim.RMSprop(policy_net.parameters(), lr=LR)
preprocess = transforms.Compose([
transforms.ToPILImage(),
transforms.RandomGrayscale(p=1),
transforms.Resize((84, 84)),
transforms.ToTensor()
])
cnt = 0
episodes_duration = []
rewards_episode = []
policy_net = policy_net.to(DEV)
target_net = target_net.to(DEV)
for episode in range(EPISODES):
state = env.reset() # image
state = preprocess(state).unsqueeze(0).to(DEV)
reward_episode = 0
for i in count():
probs = policy_net(state)
if eps_decay(cnt) > random.random():
action = torch.tensor([random.randrange(0, action_space)])
else:
action = torch.argmax(probs).view(1)
next_state, reward, done, info = env.step(action.item())
reward_episode += reward
reward = torch.tensor([reward])
next_state = preprocess(next_state).unsqueeze(0)
cnt += 1
rm.add_to_memory(
(state.to("cpu"), action.to("cpu"), reward, next_state, done))
batch = rm.get_batch_for_replay(BATCH_SIZE)
if batch is None: continue
batch = Transition(*zip(*batch))
non_final_mask = torch.tensor(tuple(
map(lambda s: s is not None, batch.next_state)),
device=DEV)
non_final_next_states = torch.cat(
[s for s in batch.next_state if s is not None]).to(DEV)
state_batch = torch.cat(batch.state).to(DEV)
action_batch = torch.cat(batch.action).to(DEV)
reward_batch = torch.cat(batch.reward).to(DEV)
d = policy_net(state_batch)
state_action_value = policy_net(state_batch).gather(
1, action_batch.reshape(-1, 1))
next_states_values = torch.zeros(BATCH_SIZE, device=DEV)
next_states_values[non_final_mask] = target_net(
non_final_next_states).max(1)[0].detach()
y = (next_states_values * GAMMA) + reward_batch
loss = F.smooth_l1_loss(state_action_value.squeeze(), y)
opt.zero_grad()
loss.backward()
for param in policy_net.parameters():
param.grad.data.clamp_(-1, 1)
opt.step()
if done:
episodes_duration.append(i + 1)
rewards_episode.append(reward_episode)
print(
f"Episode: {episode} || step: {cnt} || reward episode: {reward_episode}"
)
#plot_durations(episodes_duration)
break
if episode % TARGET_UPDATE == 0:
target_net.load_state_dict(policy_net.state_dict())
def _get_data(self):
data_dir = self.root
to_pil = transforms.ToPILImage()
to_tensor = transforms.ToTensor()
# Choose subsets that should be included into the training
training_list_img = []
training_list_labels = []
training_list_idx = []
training_list_size = []
training_out_classes = []
if self.data_case == "train":
to_tensor = transforms.Compose([
transforms.Resize((self.args.img_w, self.args.img_h)),
transforms.RandomResizedCrop(self.args.img_w,
scale=(0.7, 1.0)),
transforms.RandomHorizontalFlip(),
transforms.ColorJitter(0.3, 0.3, 0.3, 0.3),
transforms.RandomGrayscale(),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406],
[0.229, 0.224, 0.225]),
])
else:
to_tensor = transforms.Compose([
transforms.Resize((self.args.img_w, self.args.img_h)),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406],
[0.229, 0.224, 0.225]),
])
## import pdb; pdb.set_trace();
for domain in self.list_train_domains:
if self.data_case == "train":
domain_data = h5py.File(data_dir + domain + "_" + "train.hdf5",
"r")
elif self.data_case == "val":
domain_data = h5py.File(data_dir + domain + "_" + "val.hdf5",
"r")
elif self.data_case == "test":
domain_data = h5py.File(data_dir + domain + "_" + "test.hdf5",
"r")
pacs_imgs = domain_data.get("images")
# Convert labels in the range(1,7) into (0,6)
pacs_labels = np.array(domain_data.get("labels")) - 1
pacs_idx = []
print(
"Image: ",
pacs_imgs.shape,
" Labels: ",
pacs_labels.shape,
" Out Classes: ",
len(np.unique(pacs_labels)),
)
pacs_img_trans = torch.zeros((pacs_imgs.shape[0], self.args.img_c,
self.args.img_w, self.args.img_h))
for i in range(len(pacs_imgs)):
curr_img = Image.fromarray(
pacs_imgs[i, :, :, :].astype("uint8"), "RGB")
pacs_img_trans[i, :, :, :] = to_tensor(curr_img)
pacs_idx.append(i)
print("Source Domain ", domain)
training_list_img.append(pacs_img_trans)
training_list_labels.append(torch.tensor(pacs_labels).long())
training_list_idx.append(pacs_idx)
training_list_size.append(len(pacs_imgs))
training_out_classes.append(len(np.unique(pacs_labels)))
if self.data_case == "train":
num_classes = 7
for y_c in range(num_classes):
base_class_size = 0
base_class_idx = -1
for d_idx, domain in enumerate(self.list_train_domains):
class_idx = training_list_labels[d_idx] == y_c
curr_class_size = training_list_labels[d_idx][
class_idx].shape[0]
if base_class_size < curr_class_size:
base_class_size = curr_class_size
base_class_idx = d_idx
self.base_domain_size += base_class_size
print("Max Class Size: ", base_class_size, base_class_idx, y_c)
# Stack
train_imgs = torch.cat(training_list_img)
train_labels = torch.cat(training_list_labels)
train_indices = np.array(training_list_idx)
train_indices = np.hstack(train_indices)
self.training_list_size = training_list_size
print(train_imgs.shape, train_labels.shape, train_indices.shape)
print(self.training_list_size)
# Create domain labels
train_domains = torch.zeros(train_labels.size())
domain_start = 0
for idx in range(len(self.list_train_domains)):
curr_domain_size = self.training_list_size[idx]
train_domains[domain_start:domain_start + curr_domain_size] += idx
domain_start += curr_domain_size
# Shuffle everything one more time
inds = np.arange(train_labels.size()[0])
np.random.shuffle(inds)
train_imgs = train_imgs[inds]
train_labels = train_labels[inds]
train_domains = train_domains[inds].long()
train_indices = train_indices[inds]
# Convert to onehot
out_classes = training_out_classes[0]
y = torch.eye(out_classes)
train_labels = y[train_labels]
# Convert to onehot
d = torch.eye(len(self.list_train_domains))
train_domains = d[train_domains]
print(
train_imgs.shape,
train_labels.shape,
train_domains.shape,
train_indices.shape,
)
# If shape (B,H,W) change it to (B,C,H,W) with C=1
if len(train_imgs.shape) == 3:
train_imgs = train_imgs.unsqueeze(1)
return train_imgs, train_labels, train_domains, train_indices
def main():
# 随机种子
np.random.seed(666)
torch.manual_seed(666)
torch.cuda.manual_seed_all(666)
random.seed(666)
# 获取当前文件名，用于创建模型及结果文件的目录
file_name = os.path.basename(__file__).split('.')[0]
# 创建保存模型和结果的文件夹
if not os.path.exists('../data/model/%s' % file_name):
os.makedirs('../data/model/%s' % file_name)
if not os.path.exists('../data/result/%s' % file_name):
os.makedirs('../data/result/%s' % file_name)
# 创建日志文件
if not os.path.exists('../data/result/%s.txt' % file_name):
with open('../data/result/%s.txt' % file_name, 'w') as acc_file:
pass
with open('../data/result/%s.txt' % file_name, 'a') as acc_file:
acc_file.write('\n%s %s\n' % (time.strftime(
"%Y-%m-%d %H:%M:%S", time.localtime(time.time())), file_name))
# 默认使用PIL读图
def default_loader(path):
# return Image.open(path)
return Image.open(path).convert('RGB')
# 训练集图片读取
class TrainDataset(Dataset):
def __init__(self,
label_list,
transform=None,
target_transform=None,
loader=default_loader):
imgs = []
for index, row in label_list.iterrows():
imgs.append((row['img_path'], row['label']))
self.imgs = imgs
self.transform = transform
self.target_transform = target_transform
self.loader = loader
def __getitem__(self, index):
filename, label = self.imgs[index]
img = self.loader(filename)
if self.transform is not None:
img = self.transform(img)
return img, label
def __len__(self):
return len(self.imgs)
# 验证集图片读取
class ValDataset(Dataset):
def __init__(self,
label_list,
transform=None,
target_transform=None,
loader=default_loader):
imgs = []
for index, row in label_list.iterrows():
imgs.append((row['img_path'], row['label']))
self.imgs = imgs
self.transform = transform
self.target_transform = target_transform
self.loader = loader
def __getitem__(self, index):
filename, label = self.imgs[index]
img = self.loader(filename)
if self.transform is not None:
img = self.transform(img)
return img, label
def __len__(self):
return len(self.imgs)
# 测试集图片读取
class TestDataset(Dataset):
def __init__(self,
label_list,
transform=None,
target_transform=None,
loader=default_loader):
imgs = []
for index, row in label_list.iterrows():
imgs.append((row['img_path']))
self.imgs = imgs
self.transform = transform
self.target_transform = target_transform
self.loader = loader
def __getitem__(self, index):
filename = self.imgs[index]
img = self.loader(filename)
if self.transform is not None:
img = self.transform(img)
return img, filename
def __len__(self):
return len(self.imgs)
# 数据增强：在给定角度中随机进行旋转
class FixedRotation(object):
def __init__(self, angles):
self.angles = angles
def __call__(self, img):
return fixed_rotate(img, self.angles)
def fixed_rotate(img, angles):
angles = list(angles)
angles_num = len(angles)
index = random.randint(0, angles_num - 1)
return img.rotate(angles[index])
# 训练函数
def train(train_loader, model, criterion, optimizer, epoch):
batch_time = AverageMeter()
data_time = AverageMeter()
losses = AverageMeter()
acc = AverageMeter()
# switch to train mode
model.train()
end = time.time()
# 从训练集迭代器中获取训练数据
for i, (images, target) in enumerate(train_loader):
# 评估图片读取耗时
data_time.update(time.time() - end)
# 将图片和标签转化为tensor
image_var = torch.tensor(images).cuda(async=True)
label = torch.tensor(target).cuda(async=True)
# 将图片输入网络，前传，生成预测值
y_pred = model(image_var)
# 计算loss
loss = criterion(y_pred, label)
losses.update(loss.item(), images.size(0))
# 计算top1正确率
prec, PRED_COUNT = accuracy(y_pred.data, target, topk=(1, 1))
acc.update(prec, PRED_COUNT)
# 对梯度进行反向传播，使用随机梯度下降更新网络权重
optimizer.zero_grad()
loss.backward()
optimizer.step()
# 评估训练耗时
batch_time.update(time.time() - end)
end = time.time()
# 打印耗时与结果
if i % print_freq == 0:
print('Epoch: [{0}][{1}/{2}]\t'
'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
'Data {data_time.val:.3f} ({data_time.avg:.3f})\t'
'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
'Accuray {acc.val:.3f} ({acc.avg:.3f})'.format(
epoch,
i,
len(train_loader),
batch_time=batch_time,
data_time=data_time,
loss=losses,
acc=acc))
# 验证函数
def validate(val_loader, model, criterion):
batch_time = AverageMeter()
losses = AverageMeter()
acc = AverageMeter()
# switch to evaluate mode
model.eval()
end = time.time()
for i, (images, labels) in enumerate(val_loader):
image_var = torch.tensor(images).cuda(async=True)
target = torch.tensor(labels).cuda(async=True)
# 图片前传。验证和测试时不需要更新网络权重，所以使用torch.no_grad()，表示不计算梯度
with torch.no_grad():
y_pred = model(image_var)
loss = criterion(y_pred, target)
# measure accuracy and record loss
prec, PRED_COUNT = accuracy(y_pred.data, labels, topk=(1, 1))
losses.update(loss.item(), images.size(0))
acc.update(prec, PRED_COUNT)
# measure elapsed time
batch_time.update(time.time() - end)
end = time.time()
if i % print_freq == 0:
print('TrainVal: [{0}/{1}]\t'
'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
'Accuray {acc.val:.3f} ({acc.avg:.3f})'.format(
i,
len(val_loader),
batch_time=batch_time,
loss=losses,
acc=acc))
print(' * Accuray {acc.avg:.3f}'.format(acc=acc),
'(Previous Best Acc: %.3f)' % best_precision,
' * Loss {loss.avg:.3f}'.format(loss=losses),
'Previous Lowest Loss: %.3f)' % lowest_loss)
return acc.avg, losses.avg
# 测试函数
def test(test_loader, model):
csv_map = OrderedDict({'filename': [], 'probability': []})
# switch to evaluate mode
model.eval()
for i, (images, filepath) in enumerate(tqdm(test_loader)):
# bs, ncrops, c, h, w = images.size()
filepath = [os.path.basename(i) for i in filepath]
image_var = torch.tensor(images,
requires_grad=False) # for pytorch 0.4
with torch.no_grad():
y_pred = model(image_var)
# 使用softmax函数将图片预测结果转换成类别概率
smax = nn.Softmax(1)
smax_out = smax(y_pred)
# 保存图片名称与预测概率
csv_map['filename'].extend(filepath)
for output in smax_out:
prob = ';'.join([str(i) for i in output.data.tolist()])
csv_map['probability'].append(prob)
result = pd.DataFrame(csv_map)
result['probability'] = result['probability'].map(
lambda x: [float(i) for i in x.split(';')])
# 转换成提交样例中的格式
sub_filename, sub_label = [], []
for index, row in result.iterrows():
sub_filename.append(row['filename'])
pred_label = np.argmax(row['probability'])
if pred_label == 0:
sub_label.append('norm')
else:
if pred_label == 12:
pred_label = 2
if pred_label == 13:
pred_label = 11
if pred_label == 14:
pred_label = 8
if pred_label == 15:
pred_label = 10
if pred_label == 16:
pred_label = 2
if pred_label == 17:
pred_label = 10
sub_label.append('defect%d' % pred_label)
# 生成结果文件，保存在result文件夹中，可用于直接提交
submission = pd.DataFrame({
'filename': sub_filename,
'label': sub_label
})
submission.to_csv(
("../submit/submit_" +
datetime.datetime.now().strftime('%Y%m%d_%H%M%S') + ".csv"),
header=None,
index=False)
return
# 保存最新模型以及最优模型
def save_checkpoint(state,
is_best,
is_lowest_loss,
filename='../data/model/%s/checkpoint.pth.tar' %
file_name):
torch.save(state, filename)
if is_best:
shutil.copyfile(filename,
'../data/model/%s/model_best.pth.tar' % file_name)
if is_lowest_loss:
shutil.copyfile(filename,
'../data/model/%s/lowest_loss.pth.tar' % file_name)
# 用于计算精度和时间的变化
class AverageMeter(object):
"""Computes and stores the average and current value"""
def __init__(self):
self.reset()
def reset(self):
self.val = 0
self.avg = 0
self.sum = 0
self.count = 0
def update(self, val, n=1):
self.val = val
self.sum += val * n
self.count += n
self.avg = self.sum / self.count
# 学习率衰减：lr = lr / lr_decay
def adjust_learning_rate():
nonlocal lr
lr = lr / lr_decay
return optim.Adam(model.parameters(),
lr,
weight_decay=weight_decay,
amsgrad=True)
# 计算top K准确率
def accuracy(y_pred, y_actual, topk=(1, )):
"""Computes the precision@k for the specified values of k"""
final_acc = 0
maxk = max(topk)
# for prob_threshold in np.arange(0, 1, 0.01):
PRED_COUNT = y_actual.size(0)
PRED_CORRECT_COUNT = 0
prob, pred = y_pred.topk(maxk, 1, True, True)
# prob = np.where(prob > prob_threshold, prob, 0)
for j in range(pred.size(0)):
if int(y_actual[j]) == int(pred[j]):
PRED_CORRECT_COUNT += 1
if PRED_COUNT == 0:
final_acc = 0
else:
final_acc = PRED_CORRECT_COUNT / PRED_COUNT
return final_acc * 100, PRED_COUNT
# 程序主体
# 设定GPU ID
# 小数据集上，batch size不易过大。如出现out of memory，应调小batch size
batch_size = 24
# 进程数量，最好不要超过电脑最大进程数，尽量能被batch size整除。windows下报错可以改为workers=0
workers = 12
# epoch数量，分stage进行，跑完一个stage后降低学习率进入下一个stage
stage_epochs = [20, 10, 10]
# 初始学习率
lr = 5e-5
# 学习率衰减系数 (new_lr = lr / lr_decay)
lr_decay = 5
# 正则化系数
weight_decay = 1e-3
# 参数初始化
stage = 0
start_epoch = 0
total_epochs = sum(stage_epochs)
best_precision = 0
lowest_loss = 100
# 设定打印频率，即多少step打印一次，用于观察loss和acc的实时变化
# 打印结果中，括号前面为实时loss和acc，括号内部为epoch内平均loss和acc
print_freq = 1
# 验证集比例
val_ratio = 0.12
# 是否只验证，不训练
evaluate = False
# 是否从断点继续跑
resume = False
# 创建inception_v4模型
model = pretrainedmodels.__dict__['nasnetalarge'](num_classes=1000,
pretrained='imagenet')
model.last_linear = nn.Linear(4032, 18)
model = torch.nn.DataParallel(model).cuda()
# optionally resume from a checkpoint
if resume:
checkpoint_path = '../data/model/%s/checkpoint.pth.tar' % file_name
if os.path.isfile(checkpoint_path):
print("=> loading checkpoint '{}'".format(checkpoint_path))
checkpoint = torch.load(checkpoint_path)
start_epoch = checkpoint['epoch'] + 1
best_precision = checkpoint['best_precision']
lowest_loss = checkpoint['lowest_loss']
stage = checkpoint['stage']
lr = checkpoint['lr']
model.load_state_dict(checkpoint['state_dict'])
# 如果中断点恰好为转换stage的点，需要特殊处理
if start_epoch in np.cumsum(stage_epochs)[:-1]:
stage += 1
optimizer = adjust_learning_rate()
model.load_state_dict(
torch.load('../data/model/%s/model_best.pth.tar' %
file_name)['state_dict'])
print("=> loaded checkpoint (epoch {})".format(
checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(resume))
# 读取训练图片列表
all_data = pd.read_csv('../data/label.csv')
# 分离训练集和测试集，stratify参数用于分层抽样
train_data_list, val_data_list = train_test_split(
all_data,
test_size=val_ratio,
random_state=666,
stratify=all_data['label'])
# 读取测试图片列表
test_data_list = pd.read_csv('../data/test.csv')
# 图片归一化，由于采用ImageNet预训练网络，因此这里直接采用ImageNet网络的参数
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
# 训练集图片变换，输入网络的尺寸为384*384
train_data = TrainDataset(
train_data_list,
transform=transforms.Compose([
transforms.Resize((350, 350)),
transforms.ColorJitter(0.15, 0.15, 0.15, 0.075),
transforms.RandomHorizontalFlip(),
transforms.RandomGrayscale(),
# transforms.RandomRotation(20),
FixedRotation([0, 90, 180, 270]),
transforms.RandomCrop((331, 331)),
transforms.ToTensor(),
normalize,
]))
# 验证集图片变换
val_data = ValDataset(val_data_list,
transform=transforms.Compose([
transforms.Resize((350, 350)),
transforms.CenterCrop((331, 331)),
transforms.ToTensor(),
normalize,
]))
# 测试集图片变换
test_data = TestDataset(test_data_list,
transform=transforms.Compose([
transforms.Resize((350, 350)),
transforms.CenterCrop((331, 331)),
transforms.ToTensor(),
normalize,
]))
# 生成图片迭代器
train_loader = DataLoader(train_data,
batch_size=batch_size,
shuffle=True,
pin_memory=True,
num_workers=workers)
val_loader = DataLoader(val_data,
batch_size=batch_size * 2,
shuffle=False,
pin_memory=False,
num_workers=workers)
test_loader = DataLoader(test_data,
batch_size=batch_size * 2,
shuffle=False,
pin_memory=False,
num_workers=workers)
# 使用交叉熵损失函数
criterion = nn.CrossEntropyLoss().cuda()
# 优化器，使用带amsgrad的Adam
optimizer = optim.Adam(model.parameters(),
lr,
weight_decay=weight_decay,
amsgrad=True)
if evaluate:
validate(val_loader, model, criterion)
else:
# 开始训练
for epoch in range(start_epoch, total_epochs):
# train for one epoch
train(train_loader, model, criterion, optimizer, epoch)
# evaluate on validation set
precision, avg_loss = validate(val_loader, model, criterion)
# 在日志文件中记录每个epoch的精度和loss
with open('../data/result/%s.txt' % file_name, 'a') as acc_file:
acc_file.write('Epoch: %2d, Precision: %.8f, Loss: %.8f\n' %
(epoch, precision, avg_loss))
# 记录最高精度与最低loss，保存最新模型与最佳模型
is_best = precision > best_precision
is_lowest_loss = avg_loss < lowest_loss
best_precision = max(precision, best_precision)
lowest_loss = min(avg_loss, lowest_loss)
state = {
'epoch': epoch,
'state_dict': model.state_dict(),
'best_precision': best_precision,
'lowest_loss': lowest_loss,
'stage': stage,
'lr': lr,
}
save_checkpoint(state, is_best, is_lowest_loss)
# 判断是否进行下一个stage
if (epoch + 1) in np.cumsum(stage_epochs)[:-1]:
stage += 1
optimizer = adjust_learning_rate()
model.load_state_dict(
torch.load('../data/model/%s/model_best.pth.tar' %
file_name)['state_dict'])
print('Step into next stage')
with open('../data/result/%s.txt' % file_name,
'a') as acc_file:
acc_file.write(
'---------------Step into next stage----------------\n'
)
# 记录线下最佳分数
with open('../data/result/%s.txt' % file_name, 'a') as acc_file:
acc_file.write('* best acc: %.8f %s\n' %
(best_precision, os.path.basename(__file__)))
with open('../data/result/best_acc.txt', 'a') as acc_file:
acc_file.write(
'%s * best acc: %.8f %s\n' %
(time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(
time.time())), best_precision, os.path.basename(__file__)))
# 读取最佳模型，预测测试集，并生成可直接提交的结果文件
best_model = torch.load('../data/model/%s/model_best.pth.tar' % file_name)
model.load_state_dict(best_model['state_dict'])
test(test_loader=test_loader, model=model)
# 释放GPU缓存
torch.cuda.empty_cache()
def get_train_transformations(p):
if p['augmentation_strategy'] == 'standard':
# Standard augmentation strategy
return transforms.Compose([
transforms.RandomResizedCrop(
**p['augmentation_kwargs']['random_resized_crop']),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(**p['augmentation_kwargs']['normalize'])
])
elif p['augmentation_strategy'] == 'simclr':
# Augmentation strategy from the SimCLR paper
if p['train_db_name'] == 'mnist':
return transforms.Compose([
transforms.Resize(48),
transforms.RandomResizedCrop(
**p['augmentation_kwargs']['random_resized_crop']),
transforms.RandomHorizontalFlip(),
transforms.RandomApply([
transforms.ColorJitter(
**p['augmentation_kwargs']['color_jitter'])
],
p=p['augmentation_kwargs']
['color_jitter_random_apply']['p']),
transforms.RandomGrayscale(
**p['augmentation_kwargs']['random_grayscale']),
transforms.ToTensor(),
transforms.Normalize(**p['augmentation_kwargs']['normalize'])
])
else:
return transforms.Compose([
transforms.RandomResizedCrop(
**p['augmentation_kwargs']['random_resized_crop']),
transforms.RandomHorizontalFlip(),
transforms.RandomApply([
transforms.ColorJitter(
**p['augmentation_kwargs']['color_jitter'])
],
p=p['augmentation_kwargs']
['color_jitter_random_apply']['p']),
transforms.RandomGrayscale(
**p['augmentation_kwargs']['random_grayscale']),
transforms.ToTensor(),
transforms.Normalize(**p['augmentation_kwargs']['normalize'])
])
elif p['augmentation_strategy'] == 'ours':
# Augmentation strategy from our paper
return transforms.Compose([
transforms.RandomHorizontalFlip(),
transforms.RandomCrop(p['augmentation_kwargs']['crop_size'],
pad_if_needed=True),
Augment(p['augmentation_kwargs']['num_strong_augs']),
transforms.ToTensor(),
transforms.Normalize(**p['augmentation_kwargs']['normalize']),
Cutout(
n_holes=p['augmentation_kwargs']['cutout_kwargs']['n_holes'],
length=p['augmentation_kwargs']['cutout_kwargs']['length'],
random=p['augmentation_kwargs']['cutout_kwargs']['random'])
])
else:
raise ValueError('Invalid augmentation strategy {}'.format(
p['augmentation_strategy']))
'test':
transforms.Compose([
transforms.ToPILImage(),
transforms.CenterCrop(100),
#transforms.Resize(args.resize_size),
transforms.ToTensor(),
transforms.Normalize(dataset_mean, dataset_std)
])
}
## transforms for handshape annotated data
handshape_transforms = {
'train':
transforms.Compose([
transforms.RandomResizedCrop(args.random_crop_size),
#transforms.Resize(args.resize_size),
transforms.RandomGrayscale(0.2),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(dataset_mean, dataset_std)
]),
'val':
transforms.Compose([
transforms.CenterCrop(args.random_crop_size),
#transforms.Resize(args.resize_size),
transforms.ToTensor(),
transforms.Normalize(dataset_mean, dataset_std)
]),
'test':
transforms.Compose([
transforms.CenterCrop(args.random_crop_size),
#transforms.Resize(args.resize_size),
DATALOADER=dict(NUM_WORKERS=6, ),
TRAINER=dict(
FP16=dict(ENABLED=False),
NAME="SWAVRunner",
),
INPUT=dict(
AUG=dict(
TRAIN_PIPELINES=dict(
contrastive=[
("RepeatList", dict(transforms=[
("Torch_Compose", transforms.Compose([
transforms.RandomResizedCrop(224, scale=(0.14, 1.)),
transforms.RandomHorizontalFlip(p=0.5),
transforms.RandomApply([
transforms.ColorJitter(0.8, 0.8, 0.8, 0.2)], p=0.8),
transforms.RandomGrayscale(p=0.2),
])),
("RandomGaussianBlur", dict(sigma=[.1, 2.], p=0.5, mode="PIL")),
("Torch_Compose", transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(
mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])),
], repeat_times=2)),
],
multiview=[
("RepeatList", dict(transforms=[
("Torch_Compose", transforms.Compose([
transforms.RandomResizedCrop(96, scale=(0.05, 0.14)),
transforms.RandomHorizontalFlip(p=0.5),
def test_random_grayscale(self):
"""Unit tests for random grayscale transform"""
# Test Set 1: RGB -> 3 channel grayscale
random_state = random.getstate()
random.seed(42)
x_shape = [2, 2, 3]
x_np = np.random.randint(0, 256, x_shape, np.uint8)
x_pil = Image.fromarray(x_np, mode='RGB')
x_pil_2 = x_pil.convert('L')
gray_np = np.array(x_pil_2)
num_samples = 250
num_gray = 0
for _ in range(num_samples):
gray_pil_2 = transforms.RandomGrayscale(p=0.5)(x_pil)
gray_np_2 = np.array(gray_pil_2)
if np.array_equal(gray_np_2[:, :, 0], gray_np_2[:, :, 1]) and \
np.array_equal(gray_np_2[:, :, 1], gray_np_2[:, :, 2]) and \
np.array_equal(gray_np, gray_np_2[:, :, 0]):
num_gray = num_gray + 1
p_value = stats.binom_test(num_gray, num_samples, p=0.5)
random.setstate(random_state)
assert p_value > 0.0001
# Test Set 2: grayscale -> 1 channel grayscale
random_state = random.getstate()
random.seed(42)
x_shape = [2, 2, 3]
x_np = np.random.randint(0, 256, x_shape, np.uint8)
x_pil = Image.fromarray(x_np, mode='RGB')
x_pil_2 = x_pil.convert('L')
gray_np = np.array(x_pil_2)
num_samples = 250
num_gray = 0
for _ in range(num_samples):
gray_pil_3 = transforms.RandomGrayscale(p=0.5)(x_pil_2)
gray_np_3 = np.array(gray_pil_3)
if np.array_equal(gray_np, gray_np_3):
num_gray = num_gray + 1
p_value = stats.binom_test(
num_gray, num_samples,
p=1.0) # Note: grayscale is always unchanged
random.setstate(random_state)
assert p_value > 0.0001
# Test set 3: Explicit tests
x_shape = [2, 2, 3]
x_data = [0, 5, 13, 54, 135, 226, 37, 8, 234, 90, 255, 1]
x_np = np.array(x_data, dtype=np.uint8).reshape(x_shape)
x_pil = Image.fromarray(x_np, mode='RGB')
x_pil_2 = x_pil.convert('L')
gray_np = np.array(x_pil_2)
# Case 3a: RGB -> 3 channel grayscale (grayscaled)
trans2 = transforms.RandomGrayscale(p=1.0)
gray_pil_2 = trans2(x_pil)
gray_np_2 = np.array(gray_pil_2)
assert gray_pil_2.mode == 'RGB', 'mode should be RGB'
assert gray_np_2.shape == tuple(x_shape), 'should be 3 channel'
np.testing.assert_equal(gray_np_2[:, :, 0], gray_np_2[:, :, 1])
np.testing.assert_equal(gray_np_2[:, :, 1], gray_np_2[:, :, 2])
np.testing.assert_equal(gray_np, gray_np_2[:, :, 0])
# Case 3b: RGB -> 3 channel grayscale (unchanged)
trans2 = transforms.RandomGrayscale(p=0.0)
gray_pil_2 = trans2(x_pil)
gray_np_2 = np.array(gray_pil_2)
assert gray_pil_2.mode == 'RGB', 'mode should be RGB'
assert gray_np_2.shape == tuple(x_shape), 'should be 3 channel'
np.testing.assert_equal(x_np, gray_np_2)
# Case 3c: 1 channel grayscale -> 1 channel grayscale (grayscaled)
trans3 = transforms.RandomGrayscale(p=1.0)
gray_pil_3 = trans3(x_pil_2)
gray_np_3 = np.array(gray_pil_3)
assert gray_pil_3.mode == 'L', 'mode should be L'
assert gray_np_3.shape == tuple(x_shape[0:2]), 'should be 1 channel'
np.testing.assert_equal(gray_np, gray_np_3)
# Case 3d: 1 channel grayscale -> 1 channel grayscale (unchanged)
trans3 = transforms.RandomGrayscale(p=0.0)
gray_pil_3 = trans3(x_pil_2)
gray_np_3 = np.array(gray_pil_3)
assert gray_pil_3.mode == 'L', 'mode should be L'
assert gray_np_3.shape == tuple(x_shape[0:2]), 'should be 1 channel'
np.testing.assert_equal(gray_np, gray_np_3)
# Checking if RandomGrayscale can be printed as string
trans3.__repr__()
def main():
global BEST_ACC, LR_STATE
start_epoch = cfg.CLS.start_epoch # start from epoch 0 or last checkpoint epoch
# Create ckpt folder
if not os.path.isdir(cfg.CLS.ckpt):
mkdir_p(cfg.CLS.ckpt)
if args.cfg_file is not None and not cfg.CLS.evaluate:
shutil.copyfile(
args.cfg_file,
os.path.join(cfg.CLS.ckpt,
args.cfg_file.split('/')[-1]))
# Dataset and Loader
normalize = transforms.Normalize(mean=cfg.pixel_mean, std=cfg.pixel_std)
train_aug = [
transforms.RandomResizedCrop(cfg.CLS.crop_size),
transforms.RandomHorizontalFlip()
]
if len(cfg.CLS.rotation) > 0:
train_aug.append(transforms.RandomRotation(cfg.CLS.rotation))
if len(cfg.CLS.pixel_jitter) > 0:
train_aug.append(RandomPixelJitter(cfg.CLS.pixel_jitter))
if cfg.CLS.grayscale > 0:
train_aug.append(transforms.RandomGrayscale(cfg.CLS.grayscale))
train_aug.append(transforms.ToTensor())
train_aug.append(normalize)
traindir = os.path.join(cfg.CLS.data_root, cfg.CLS.train_folder)
train_loader = torch.utils.data.DataLoader(datasets.ImageFolder(
traindir, transforms.Compose(train_aug)),
batch_size=cfg.CLS.train_batch,
shuffle=True,
num_workers=cfg.workers,
pin_memory=True)
if cfg.CLS.validate or cfg.CLS.evaluate:
valdir = os.path.join(cfg.CLS.data_root, cfg.CLS.val_folder)
val_loader = torch.utils.data.DataLoader(datasets.ImageFolder(
valdir,
transforms.Compose([
transforms.Resize(cfg.CLS.base_size),
transforms.CenterCrop(cfg.CLS.crop_size),
transforms.ToTensor(),
normalize,
])),
batch_size=cfg.CLS.test_batch,
shuffle=False,
num_workers=cfg.workers,
pin_memory=True)
# Create model
model = models.__dict__[cfg.CLS.arch]()
print(model)
# Calculate FLOPs & Param
n_flops, n_convops, n_params = measure_model(model, cfg.CLS.crop_size,
cfg.CLS.crop_size)
print('==> FLOPs: {:.4f}M, Conv_FLOPs: {:.4f}M, Params: {:.4f}M'.format(
n_flops / 1e6, n_convops / 1e6, n_params / 1e6))
del model
model = models.__dict__[cfg.CLS.arch]()
# Load pre-train model
if cfg.CLS.pretrained:
print("==> Using pre-trained model '{}'".format(cfg.CLS.pretrained))
pretrained_dict = torch.load(cfg.CLS.pretrained)
try:
pretrained_dict = pretrained_dict['state_dict']
except:
pretrained_dict = pretrained_dict
model_dict = model.state_dict()
updated_dict, match_layers, mismatch_layers = weight_filler(
pretrained_dict, model_dict)
model_dict.update(updated_dict)
model.load_state_dict(model_dict)
else:
print("==> Creating model '{}'".format(cfg.CLS.arch))
# Define loss function (criterion) and optimizer
criterion = nn.CrossEntropyLoss().cuda()
if cfg.CLS.pretrained:
def param_filter(param):
return param[1]
new_params = map(
param_filter,
filter(lambda p: p[0] in mismatch_layers,
model.named_parameters()))
base_params = map(
param_filter,
filter(lambda p: p[0] in match_layers, model.named_parameters()))
model_params = [{
'params': base_params
}, {
'params': new_params,
'lr': cfg.CLS.base_lr * 10
}]
else:
model_params = model.parameters()
model = torch.nn.DataParallel(model).cuda()
cudnn.benchmark = True
optimizer = optim.SGD(model_params,
lr=cfg.CLS.base_lr,
momentum=cfg.CLS.momentum,
weight_decay=cfg.CLS.weight_decay)
# Evaluate model
if cfg.CLS.evaluate:
print('\n==> Evaluation only')
test_loss, test_top1, test_top5 = test(val_loader, model, criterion,
start_epoch, USE_CUDA)
print(
'==> Test Loss: {:.8f} | Test_top1: {:.4f}% | Test_top5: {:.4f}%'.
format(test_loss, test_top1, test_top5))
return
# Resume training
title = 'Pytorch-CLS-' + cfg.CLS.arch
if cfg.CLS.resume:
# Load checkpoint.
print("==> Resuming from checkpoint '{}'".format(cfg.CLS.resume))
assert os.path.isfile(
cfg.CLS.resume), 'Error: no checkpoint directory found!'
checkpoint = torch.load(cfg.CLS.resume)
BEST_ACC = checkpoint['best_acc']
start_epoch = checkpoint['epoch']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
logger = Logger(os.path.join(cfg.CLS.ckpt, 'log.txt'),
title=title,
resume=True)
else:
logger = Logger(os.path.join(cfg.CLS.ckpt, 'log.txt'), title=title)
logger.set_names([
'Learning Rate', 'Train Loss', 'Valid Loss', 'Train Acc.',
'Valid Acc.'
])
# Train and val
for epoch in range(start_epoch, cfg.CLS.epochs):
print('\nEpoch: [{}/{}] | LR: {:.8f}'.format(epoch + 1, cfg.CLS.epochs,
LR_STATE))
train_loss, train_acc = mixup_train(train_loader, model, criterion,
optimizer, epoch, USE_CUDA)
if cfg.CLS.validate:
test_loss, test_top1, test_top5 = test(val_loader, model,
criterion, epoch, USE_CUDA)
else:
test_loss, test_top1, test_top5 = 0.0, 0.0, 0.0
# Append logger file
logger.append([LR_STATE, train_loss, test_loss, train_acc, test_top1])
# Save model
save_checkpoint(model, optimizer, test_top1, epoch)
# Draw curve
try:
draw_curve(cfg.CLS.arch, cfg.CLS.ckpt)
print('==> Success saving log curve...')
except:
print('==> Saving log curve error...')
logger.close()
try:
savefig(os.path.join(cfg.CLS.ckpt, 'log.eps'))
shutil.copyfile(
os.path.join(cfg.CLS.ckpt, 'log.txt'),
os.path.join(
cfg.CLS.ckpt, 'log{}.txt'.format(
datetime.datetime.now().strftime('%Y%m%d%H%M%S'))))
except:
print('copy log error.')
print('==> Training Done!')
print('==> Best acc: {:.4f}%'.format(BEST_ACC))
