def __init__(self):
# GPUが利用可能であればGPUを使用，利用可能でなければCPUを使用
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# モデルを定義しself.deviceのメモリに送る
self.model = PretrainedResNet(512).to(self.device)
# 学習済み重みの読み込み
self.model.load_state_dict(torch.load("./network/model/arcface_SGD"))
# 損失関数を定義しself.deviceのメモリに送る
self.loss = losses.ArcFaceLoss(margin=14.3, scale=64, num_classes=40, embedding_size=512).to(self.device)
# 損失関数の学習済み重みの読み込み(ArcFaceLossの内部には学習可能な部分があるため，その部分の重みを読み込む)
self.loss.load_state_dict(torch.load("./network/model/arcface_SGD_loss"))
# モデルを評価モード(重みを固定するモード)に変更
self.model.eval()
# 画像の前処理方法を指定
# 144×144にリサイズ
# Torch Tensor(pytorchで使用可能な配列)に変換
# RGB各チャネルの画素値を平均0.5，標準偏差0.5の正規分布に従うように変換(およそ0~1の範囲になる)
self.transform = transforms.Compose([transforms.Resize((144, 144)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])])
# モデルの最終出力は損失関数から得る必要がある(モデルの出力は512次元のベクトル，得たい出力は40次元のベクトル)
self.logit_getter = LogitGetter(self.loss)
# argumentとクラス名の対応表
# モデル出力は入力が各クラスである確率として得られるため，対応する配列を作っておく必要がある
self.pred_class = ("新垣結衣", "新木優子", "新田真剣佑", "綾瀬はるか", "ディーンフジオカ", "深田恭子", "浜辺美波",
"橋本環奈", "比嘉愛未", "平野紫耀", "広瀬すず", "本田翼", "井川遥", "石原さとみ", "桐谷美玲",
"北川景子", "松坂桃李", "三浦春馬", "向井理", "長澤まさみ", "中川大志", "中条あやみ", "中村倫也",
"小栗旬", "岡田将生", "坂口健太郎", "佐々木希","佐藤健", "沢尻エリカ", "瀬戸康史", "白石麻衣",
"柴咲コウ", "菅田将暉", "竹野内豊", "玉木宏", "山本美月", "山下智久", "山崎賢人", "横浜流星", "吉沢亮")
def __init__(self, nb_classes, sz_embed, margin=28.6, scale=220, **kwargs):
super(ArcFaceLoss, self).__init__()
self.nb_classes = nb_classes
self.sz_embed = sz_embed
self.margin = margin
self.scale = scale
self.loss_func = losses.ArcFaceLoss(num_classes=self.nb_classes,
embedding_size=self.sz_embed,
margin=self.margin,
scale=self.scale)
def arcface_loss(trial,
num_classes,
embedding_size,
margin_range=(0.0, 90.0),
scale_range=(0.01, 100),
**kwargs):
margin = trial.suggest_uniform("arc_margin", *margin_range)
scale = trial.suggest_uniform("scale", *scale_range)
loss = losses.ArcFaceLoss(num_classes=num_classes,
embedding_size=embedding_size,
margin=margin,
scale=scale,
**sample_regularizer(trial))
return {"loss": loss, "param": True}
def __init__(self, args):
super().__init__()
#ckpt = '../../../Masterproef/thesis_code/recognition/wandb/run-20210409_112741-28fdpx5s/files/thesis/28fdpx5s/checkpoints/epoch=299-step=18899.ckpt'
self.model = torchvision.models.resnet18(pretrained=True)
self.model.fc = nn.Linear(
512, 2378, True) #Change fully connected layer to 2379 output
self.args = args
self.extractor = torch.nn.Sequential(
OrderedDict(list(self.model.named_children())[:-1]), )
self.classifier = torch.nn.Sequential(
OrderedDict(list(self.model.named_children())[-1:]))
#Enkel CrossEntropy loss werkt momenteel
if self.args.loss == 'CrossEntropy':
self.loss = torch.nn.CrossEntropyLoss()
self.loss_requires_classifier = True
elif self.args.loss == 'ArcFace':
self.loss = losses.ArcFaceLoss(
margin=0.5,
embedding_size=self.classifier.fc.in_features,
num_classes=self.classifier.fc.out_features)
self.loss_requires_classifier = False
elif self.args.loss == 'ContrastiveLoss':
self.loss = losses.ContrastiveLoss(pos_margin=0, neg_margin=1)
self.loss_requires_classifier = False
#sampler toevoegen!!!!
elif self.args.loss == 'TripletMargin':
self.loss = losses.TripletMarginLoss(margin=0.1)
self.loss_requires_classifier = False
elif self.args.loss == 'CircleLoss':
self.loss = losses.CircleLoss(m=0.4, gamma=80)
self.loss_requires_classifier = False
else:
raise ValueError(f'Unsupported loss: {self.args.loss}')
])
# Set the datasets
train_dataset = datasets.CIFAR100(root="CIFAR100_Dataset",
train=True,
transform=train_transform,
download=True)
val_dataset = datasets.CIFAR100(root="CIFAR100_Dataset",
train=False,
transform=val_transform,
download=True)
# Set the loss functions. loss0 will be applied to the first embedder, loss1 to the second embedder etc.
loss0 = losses.TripletMarginLoss(margin=0.01)
loss1 = losses.MultiSimilarityLoss(alpha=0.1, beta=40, base=0.5)
loss2 = losses.ArcFaceLoss(margin=30, num_classes=100,
embedding_size=64).to(device)
# Set the mining functions. In this example we'll apply mining to the 2nd and 3rd cascaded outputs.
miner1 = miners.MultiSimilarityMiner(epsilon=0.1)
miner2 = miners.HDCMiner(filter_percentage=0.25)
# Set the dataloader sampler
sampler = samplers.MPerClassSampler(train_dataset.targets, m=4)
# Set other training parameters
batch_size = 32
num_epochs = 2
iterations_per_epoch = 100
# Package the above stuff into dictionaries.
models = {"trunk": trunk, "embedder": embedder}
def train_eval(args, train_data, dev_data):
logger = logging.getLogger("main")
# Create dataset & dataloader
trans = [
transforms.Resize((224, 224)),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
]
trans = transforms.Compose(trans)
train_dataset, train_char_idx = \
create_dataset(args.root, train_data, trans)
train_sampler = MetricBatchSampler(train_dataset,
train_char_idx,
n_max_per_char=args.n_max_per_char,
n_batch_size=args.n_batch_size,
n_random=args.n_random)
train_dataloader = DataLoader(train_dataset,
batch_sampler=train_sampler,
collate_fn=collate_fn)
# number of batches given to trainer
n_batch = int(len(train_dataloader))
eval_train_dataloaders = \
prepare_evaluation_dataloaders(args, args.eval_split*3, train_data, trans)
eval_dev_dataloaders = \
prepare_evaluation_dataloaders(args, args.eval_split, dev_data, trans)
# Construct model & optimizer
device = "cpu" if args.gpu < 0 else "cuda:{}".format(args.gpu)
trunk, model = create_models(args.emb_dim, args.dropout)
trunk.to(device)
model.to(device)
if args.metric_loss == "triplet":
loss_func = losses.TripletMarginLoss(
margin=args.margin,
normalize_embeddings=args.normalize,
smooth_loss=args.smooth)
elif args.metric_loss == "arcface":
loss_func = losses.ArcFaceLoss(margin=args.margin,
num_classes=len(train_data),
embedding_size=args.emb_dim)
loss_func.to(device)
if args.optimizer == "SGD":
trunk_optimizer = torch.optim.SGD(trunk.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.decay)
model_optimizer = torch.optim.SGD(model.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.decay)
optimizers = {
"trunk_optimizer": trunk_optimizer,
"embedder_optimizer": model_optimizer
}
if args.metric_loss == "arcface":
loss_optimizer = torch.optim.SGD(loss_func.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.decay)
optimizers["loss_optimizer"] = loss_optimizer
elif args.optimizer == "Adam":
trunk_optimizer = torch.optim.Adam(trunk.parameters(),
lr=args.lr,
weight_decay=args.decay)
model_optimizer = torch.optim.Adam(model.parameters(),
lr=args.lr,
weight_decay=args.decay)
optimizers = {
"trunk_optimizer": trunk_optimizer,
"embedder_optimizer": model_optimizer
}
if args.metric_loss == "arcface":
loss_optimizer = torch.optim.Adam(loss_func.parameters(),
lr=args.lr,
weight_decay=args.decay)
optimizers["loss_optimizer"] = loss_optimizer
else:
raise NotImplementedError
def lr_func(step):
if step < args.warmup:
return (step + 1) / args.warmup
else:
steps_decay = step // args.decay_freq
return 1 / args.decay_factor**steps_decay
trunk_scheduler = torch.optim.lr_scheduler.LambdaLR(
trunk_optimizer, lr_func)
model_scheduler = torch.optim.lr_scheduler.LambdaLR(
model_optimizer, lr_func)
schedulers = {
"trunk_scheduler": trunk_scheduler,
"model_scheduler": model_scheduler
}
if args.miner == "none":
mining_funcs = {}
elif args.miner == "batch-hard":
mining_funcs = {
"post_gradient_miner": miners.BatchHardMiner(use_similarity=True)
}
best_dev_eer = 1.0
i_epoch = 0
def end_of_epoch_hook(trainer):
nonlocal i_epoch, best_dev_eer
logger.info(f"EPOCH\t{i_epoch}")
if i_epoch % args.eval_freq == 0:
train_eer, train_eer_std = evaluate(args, trainer.models["trunk"],
trainer.models["embedder"],
eval_train_dataloaders)
dev_eer, dev_eer_std = evaluate(args, trainer.models["trunk"],
trainer.models["embedder"],
eval_dev_dataloaders)
logger.info("Eval EER (mean, std):\t{}\t{}".format(
train_eer, train_eer_std))
logger.info("Eval EER (mean, std):\t{}\t{}".format(
dev_eer, dev_eer_std))
if dev_eer < best_dev_eer:
logger.info("New best model!")
best_dev_eer = dev_eer
i_epoch += 1
def end_of_iteration_hook(trainer):
for scheduler in schedulers.values():
scheduler.step()
trainer = trainers.MetricLossOnly(
models={
"trunk": trunk,
"embedder": model
},
optimizers=optimizers,
batch_size=None,
loss_funcs={"metric_loss": loss_func},
mining_funcs=mining_funcs,
iterations_per_epoch=n_batch,
dataset=train_dataset,
data_device=None,
loss_weights=None,
sampler=train_sampler,
collate_fn=collate_fn,
lr_schedulers=None,
end_of_epoch_hook=end_of_epoch_hook,
end_of_iteration_hook=end_of_iteration_hook,
dataloader_num_workers=1)
trainer.train(num_epochs=args.epoch)
if args.save_model:
save_models = {
"trunk": trainer.models["trunk"].state_dict(),
"embedder": trainer.models["embedder"].state_dict(),
"args": [args.emb_dim]
}
torch.save(save_models, f"model/{args.suffix}.mdl")
return best_dev_eer
def train_model(model, model_test, criterion, optimizer, scheduler, num_epochs=25):
since = time.time()
# best_model_wts = model.state_dict()
# best_acc = 0.0
warm_up = 0.1 # We start from the 0.1*lrRate
warm_iteration = round(dataset_sizes['satellite'] / opt.batchsize) * opt.warm_epoch # first 5 epoch
if opt.arcface:
criterion_arcface = losses.ArcFaceLoss(num_classes=opt.nclasses, embedding_size=512)
if opt.cosface:
criterion_cosface = losses.CosFaceLoss(num_classes=opt.nclasses, embedding_size=512)
if opt.circle:
criterion_circle = CircleLoss(m=0.25, gamma=32) # gamma = 64 may lead to a better result.
if opt.triplet:
miner = miners.MultiSimilarityMiner()
criterion_triplet = losses.TripletMarginLoss(margin=0.3)
if opt.lifted:
criterion_lifted = losses.GeneralizedLiftedStructureLoss(neg_margin=1, pos_margin=0)
if opt.contrast:
criterion_contrast = losses.ContrastiveLoss(pos_margin=0, neg_margin=1)
if opt.sphere:
criterion_sphere = losses.SphereFaceLoss(num_classes=opt.nclasses, embedding_size=512, margin=4)
for epoch in range(num_epochs - start_epoch):
epoch = epoch + start_epoch
print('Epoch {}/{}'.format(epoch, num_epochs - 1))
print('-' * 10)
# Each epoch has a training and validation phase
for phase in ['train']:
if phase == 'train':
model.train(True) # Set model to training mode
else:
model.train(False) # Set model to evaluate mode
running_loss = 0.0
running_corrects = 0.0
running_corrects2 = 0.0
running_corrects3 = 0.0
# Iterate over data.
for data, data2, data3, data4 in zip(dataloaders['satellite'], dataloaders['street'], dataloaders['drone'],
dataloaders['google']):
# get the inputs
inputs, labels = data
inputs2, labels2 = data2
inputs3, labels3 = data3
inputs4, labels4 = data4
now_batch_size, c, h, w = inputs.shape
if now_batch_size < opt.batchsize: # skip the last batch
continue
if use_gpu:
inputs = Variable(inputs.cuda().detach())
inputs2 = Variable(inputs2.cuda().detach())
inputs3 = Variable(inputs3.cuda().detach())
labels = Variable(labels.cuda().detach())
labels2 = Variable(labels2.cuda().detach())
labels3 = Variable(labels3.cuda().detach())
if opt.extra_Google:
inputs4 = Variable(inputs4.cuda().detach())
labels4 = Variable(labels4.cuda().detach())
else:
inputs, labels = Variable(inputs), Variable(labels)
# zero the parameter gradients
optimizer.zero_grad()
# forward
if phase == 'val':
with torch.no_grad():
outputs, outputs2 = model(inputs, inputs2)
else:
if opt.views == 2:
outputs, outputs2 = model(inputs, inputs2)
elif opt.views == 3:
if opt.extra_Google:
outputs, outputs2, outputs3, outputs4 = model(inputs, inputs2, inputs3, inputs4)
else:
outputs, outputs2, outputs3 = model(inputs, inputs2, inputs3)
return_feature = opt.arcface or opt.cosface or opt.circle or opt.triplet or opt.contrast or opt.lifted or opt.sphere
if opt.views == 2:
_, preds = torch.max(outputs.data, 1)
_, preds2 = torch.max(outputs2.data, 1)
loss = criterion(outputs, labels) + criterion(outputs2, labels2)
elif opt.views == 3:
if return_feature:
logits, ff = outputs
logits2, ff2 = outputs2
logits3, ff3 = outputs3
fnorm = torch.norm(ff, p=2, dim=1, keepdim=True)
fnorm2 = torch.norm(ff2, p=2, dim=1, keepdim=True)
fnorm3 = torch.norm(ff3, p=2, dim=1, keepdim=True)
ff = ff.div(fnorm.expand_as(ff)) # 8*512,tensor
ff2 = ff2.div(fnorm2.expand_as(ff2))
ff3 = ff3.div(fnorm3.expand_as(ff3))
loss = criterion(logits, labels) + criterion(logits2, labels2) + criterion(logits3, labels3)
_, preds = torch.max(logits.data, 1)
_, preds2 = torch.max(logits2.data, 1)
_, preds3 = torch.max(logits3.data, 1)
# Multiple perspectives are combined to calculate losses, please join ''--loss_merge'' in run.sh
if opt.loss_merge:
ff_all = torch.cat((ff, ff2, ff3), dim=0)
labels_all = torch.cat((labels, labels2, labels3), dim=0)
if opt.extra_Google:
logits4, ff4 = outputs4
fnorm4 = torch.norm(ff4, p=2, dim=1, keepdim=True)
ff4 = ff4.div(fnorm4.expand_as(ff4))
loss = criterion(logits, labels) + criterion(logits2, labels2) + criterion(logits3, labels3) +criterion(logits4, labels4)
if opt.loss_merge:
ff_all = torch.cat((ff_all, ff4), dim=0)
labels_all = torch.cat((labels_all, labels4), dim=0)
if opt.arcface:
if opt.loss_merge:
loss += criterion_arcface(ff_all, labels_all)
else:
loss += criterion_arcface(ff, labels) + criterion_arcface(ff2, labels2) + criterion_arcface(ff3, labels3) # /now_batch_size
if opt.extra_Google:
loss += criterion_arcface(ff4, labels4) # /now_batch_size
if opt.cosface:
if opt.loss_merge:
loss += criterion_cosface(ff_all, labels_all)
else:
loss += criterion_cosface(ff, labels) + criterion_cosface(ff2, labels2) + criterion_cosface(ff3, labels3) # /now_batch_size
if opt.extra_Google:
loss += criterion_cosface(ff4, labels4) # /now_batch_size
if opt.circle:
if opt.loss_merge:
loss += criterion_circle(*convert_label_to_similarity(ff_all, labels_all)) / now_batch_size
else:
loss += criterion_circle(*convert_label_to_similarity(ff, labels)) / now_batch_size + criterion_circle(*convert_label_to_similarity(ff2, labels2)) / now_batch_size + criterion_circle(*convert_label_to_similarity(ff3, labels3)) / now_batch_size
if opt.extra_Google:
loss += criterion_circle(*convert_label_to_similarity(ff4, labels4)) / now_batch_size
if opt.triplet:
if opt.loss_merge:
hard_pairs_all = miner(ff_all, labels_all)
loss += criterion_triplet(ff_all, labels_all, hard_pairs_all)
else:
hard_pairs = miner(ff, labels)
hard_pairs2 = miner(ff2, labels2)
hard_pairs3 = miner(ff3, labels3)
loss += criterion_triplet(ff, labels, hard_pairs) + criterion_triplet(ff2, labels2, hard_pairs2) + criterion_triplet(ff3, labels3, hard_pairs3)# /now_batch_size
if opt.extra_Google:
hard_pairs4 = miner(ff4, labels4)
loss += criterion_triplet(ff4, labels4, hard_pairs4)
if opt.lifted:
if opt.loss_merge:
loss += criterion_lifted(ff_all, labels_all)
else:
loss += criterion_lifted(ff, labels) + criterion_lifted(ff2, labels2) + criterion_lifted(ff3, labels3) # /now_batch_size
if opt.extra_Google:
loss += criterion_lifted(ff4, labels4)
if opt.contrast:
if opt.loss_merge:
loss += criterion_contrast(ff_all, labels_all)
else:
loss += criterion_contrast(ff, labels) + criterion_contrast(ff2,labels2) + criterion_contrast(ff3, labels3) # /now_batch_size
if opt.extra_Google:
loss += criterion_contrast(ff4, labels4)
if opt.sphere:
if opt.loss_merge:
loss += criterion_sphere(ff_all, labels_all) / now_batch_size
else:
loss += criterion_sphere(ff, labels) / now_batch_size + criterion_sphere(ff2, labels2) / now_batch_size + criterion_sphere(ff3, labels3) / now_batch_size
if opt.extra_Google:
loss += criterion_sphere(ff4, labels4)
else:
_, preds = torch.max(outputs.data, 1)
_, preds2 = torch.max(outputs2.data, 1)
_, preds3 = torch.max(outputs3.data, 1)
if opt.loss_merge:
outputs_all = torch.cat((outputs, outputs2, outputs3), dim=0)
labels_all = torch.cat((labels, labels2, labels3), dim=0)
if opt.extra_Google:
outputs_all = torch.cat((outputs_all, outputs4), dim=0)
labels_all = torch.cat((labels_all, labels4), dim=0)
loss = 4*criterion(outputs_all, labels_all)
else:
loss = criterion(outputs, labels) + criterion(outputs2, labels2) + criterion(outputs3, labels3)
if opt.extra_Google:
loss += criterion(outputs4, labels4)
# backward + optimize only if in training phase
if epoch < opt.warm_epoch and phase == 'train':
warm_up = min(1.0, warm_up + 0.9 / warm_iteration)
loss *= warm_up
if phase == 'train':
if fp16: # we use optimier to backward loss
with amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()
else:
loss.backward()
optimizer.step()
##########
if opt.moving_avg < 1.0:
update_average(model_test, model, opt.moving_avg)
# statistics
if int(version[0]) > 0 or int(version[2]) > 3: # for the new version like 0.4.0, 0.5.0 and 1.0.0
running_loss += loss.item() * now_batch_size
else: # for the old version like 0.3.0 and 0.3.1
running_loss += loss.data[0] * now_batch_size
running_corrects += float(torch.sum(preds == labels.data))
running_corrects2 += float(torch.sum(preds2 == labels2.data))
if opt.views == 3:
running_corrects3 += float(torch.sum(preds3 == labels3.data))
epoch_loss = running_loss / dataset_sizes['satellite']
epoch_acc = running_corrects / dataset_sizes['satellite']
epoch_acc2 = running_corrects2 / dataset_sizes['satellite']
if opt.views == 2:
print('{} Loss: {:.4f} Satellite_Acc: {:.4f} Street_Acc: {:.4f}'.format(phase, epoch_loss, epoch_acc,
epoch_acc2))
elif opt.views == 3:
epoch_acc3 = running_corrects3 / dataset_sizes['satellite']
print('{} Loss: {:.4f} Satellite_Acc: {:.4f} Street_Acc: {:.4f} Drone_Acc: {:.4f}'.format(phase,
epoch_loss,
epoch_acc,
epoch_acc2,
epoch_acc3))
y_loss[phase].append(epoch_loss)
y_err[phase].append(1.0 - epoch_acc)
# deep copy the model
if phase == 'train':
scheduler.step()
last_model_wts = model.state_dict()
if epoch % 20 == 19:
save_network(model, opt.name, epoch)
# draw_curve(epoch)
time_elapsed = time.time() - since
print('Training complete in {:.0f}m {:.0f}s'.format(
time_elapsed // 60, time_elapsed % 60))
print()
time_elapsed = time.time() - since
print('Training complete in {:.0f}m {:.0f}s'.format(
time_elapsed // 60, time_elapsed % 60))
# print('Best val Acc: {:4f}'.format(best_acc))
# save_network(model_test, opt.name+'adapt', epoch)
return model
def main():
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print('Running on device: {}'.format(device))
# Data transformations
trans_train = transforms.Compose([
transforms.RandomApply(transforms=[
transforms.AutoAugment(transforms.AutoAugmentPolicy.IMAGENET),
# transforms.RandomPerspective(distortion_scale=0.6, p=1.0),
transforms.RandomRotation(degrees=(0, 180)),
transforms.RandomHorizontalFlip(),
]),
np.float32,
transforms.ToTensor(),
fixed_image_standardization,
])
trans_val = transforms.Compose([
# transforms.CenterCrop(120),
np.float32,
transforms.ToTensor(),
fixed_image_standardization,
])
train_dataset = datasets.ImageFolder(os.path.join(data_dir,
"train_aligned"),
transform=trans_train)
val_dataset = datasets.ImageFolder(os.path.join(data_dir, "val_aligned"),
transform=trans_val)
# Prepare the model
model = InceptionResnetV1(classify=False,
pretrained="vggface2",
dropout_prob=0.5).to(device)
# for param in list(model.parameters())[:-8]:
# param.requires_grad = False
trunk_optimizer = torch.optim.SGD(model.parameters(), lr=LR)
# Set the loss function
loss = losses.ArcFaceLoss(len(train_dataset.classes), 512)
# Package the above stuff into dictionaries.
models = {"trunk": model}
optimizers = {"trunk_optimizer": trunk_optimizer}
loss_funcs = {"metric_loss": loss}
mining_funcs = {}
lr_scheduler = {
"trunk_scheduler_by_plateau":
torch.optim.lr_scheduler.ReduceLROnPlateau(trunk_optimizer)
}
# Create the tester
record_keeper, _, _ = logging_presets.get_record_keeper(
"logs", "tensorboard")
hooks = logging_presets.get_hook_container(record_keeper)
dataset_dict = {"val": val_dataset, "train": train_dataset}
model_folder = "training_saved_models"
def visualizer_hook(umapper, umap_embeddings, labels, split_name, keyname,
*args):
logging.info("UMAP plot for the {} split and label set {}".format(
split_name, keyname))
label_set = np.unique(labels)
num_classes = len(label_set)
fig = plt.figure(figsize=(8, 7))
plt.gca().set_prop_cycle(
cycler("color", [
plt.cm.nipy_spectral(i)
for i in np.linspace(0, 0.9, num_classes)
]))
for i in range(num_classes):
idx = labels == label_set[i]
plt.plot(umap_embeddings[idx, 0],
umap_embeddings[idx, 1],
".",
markersize=1)
plt.show()
tester = testers.GlobalEmbeddingSpaceTester(
end_of_testing_hook=hooks.end_of_testing_hook,
dataloader_num_workers=4,
accuracy_calculator=AccuracyCalculator(
include=['mean_average_precision_at_r'], k="max_bin_count"))
end_of_epoch_hook = hooks.end_of_epoch_hook(tester,
dataset_dict,
model_folder,
splits_to_eval=[('val',
['train'])])
# Create the trainer
trainer = trainers.MetricLossOnly(
models,
optimizers,
batch_size,
loss_funcs,
mining_funcs,
train_dataset,
lr_schedulers=lr_scheduler,
dataloader_num_workers=8,
end_of_iteration_hook=hooks.end_of_iteration_hook,
end_of_epoch_hook=end_of_epoch_hook)
trainer.train(num_epochs=num_epochs)
def __init__(self, setting_csv_path, index):
self.parameters_dict = read_parameters(
setting_csv_path, index) # 全ハイパーパラメータが保存されたディクショナリ
self.model_name = self.parameters_dict["model_name"] # モデル名
log_dir_name = self.model_name + "_epochs" + self.parameters_dict["epochs"] \
+ "_batch_size" + self.parameters_dict["batch_size"] \
+ "_lr" + self.parameters_dict["learning_rate"] \
+ "_margin" + self.parameters_dict["margin"] \
+ "_scale" + self.parameters_dict["scale"] # ログを保存するフォルダ名
self.log_path = os.path.join(self.parameters_dict["base_log_path"],
log_dir_name) # ログの保存先
batch_size = int(self.parameters_dict["batch_size"]) # バッチサイズ
learning_rate = float(self.parameters_dict["learning_rate"]) # 学習率
momentum = float(self.parameters_dict["momentum"]) # 慣性項(SGD使用時のみ使用)
weight_decay = float(
self.parameters_dict["weight_decay"]) # 重み減衰(SGD使用時のみ使用)
img_size = (int(self.parameters_dict["width"]),
int(self.parameters_dict["height"])) # 画像サイズ
self.logger = Logger(self.log_path) # ログ書き込みを行うLoggerクラスの宣言
num_class = int(self.parameters_dict["num_class"]) # クラス数
step_size = int(self.parameters_dict["step_size"])
self.device = torch.device("cuda" if torch.cuda.is_available() else
"cpu") # GPUが利用可能であればGPUを利用
#self.model = VGG_based(num_class).to(self.device) # ネットワークを定義(VGGベース)
#self.model = ResNet_based(num_class).to(self.device) # ネットワークを定義(ResNetベース)
self.model = PretrainedResNet(num_class).to(
self.device) # ネットワークを定義(学習済みResNetベース)
#学習済み重みファイルがあるか確認しあれば読み込み
if os.path.isfile(os.path.join(self.log_path, self.model_name)):
print("Trained weight file exists")
self.model.load_state_dict(
torch.load(os.path.join(self.log_path, self.model_name)))
#CNN部分の最適化手法の定義
#ArcFaceLoss
#簡単のためpytorch_metric_learningからimportして読み込み
#margin : クラス間の分離を行う際の最少距離(cosine類似度による距離学習を行うためmarginはθを示す)
#scale : クラスをどの程度の大きさに収めるか
#num_classes : ArcFaceLossにはMLPが含まれるためMLPのパラメータとして入力
#embedding_size : 同上
self.loss = losses.ArcFaceLoss(
margin=float(self.parameters_dict["margin"]),
scale=int(self.parameters_dict["scale"]),
num_classes=num_class,
embedding_size=512).to(self.device)
#self.optimizer = torch.optim.Adam(list(self.model.parameters()) + list(self.loss.parameters()), lr=self.learning_rate) # PMLのArcFaceLossにはMLPが含まれている(Trainable)なのでモデルパラメータとlossに含まれるモデルパラメータを最適化
self.optimizer = torch.optim.SGD(
list(self.model.parameters()) + list(self.loss.parameters()),
lr=learning_rate,
momentum=momentum,
weight_decay=weight_decay,
nesterov=True
) # PMLのArcFaceLossにはMLPが含まれている(Trainable)なのでモデルパラメータとlossに含まれるモデルパラメータを最適化
self.scheduler = torch.optim.lr_scheduler.StepLR(
optimizer=self.optimizer, step_size=step_size, gamma=0.1)
#バッチ読み込みをいくつのスレッドに並列化するか指定
#パラメータ辞書に"-1"と登録されていればCPUのコア数を読み取って指定
num_workers = 0
if int(self.parameters_dict["num_workers"]) == -1:
print("set num_workers to number of cpu cores :", os.cpu_count())
num_workers = os.cpu_count()
else:
num_workers = int(self.parameters_dict["num_workers"])
#データローダーの定義
#data_path : データの保存先
#batch_size : バッチサイズ
#img_size : 画像サイズ(タプルで指定)
#train_ratio : 全データ中学習に使用するデータの割合
self.data_loader = DataLoader(
data_path=self.parameters_dict["data_path"],
batch_size=int(self.parameters_dict["batch_size"]),
img_size=img_size,
train_ratio=float(self.parameters_dict["train_ratio"]),
num_workers=num_workers,
pin_memory=str_to_bool(self.parameters_dict["pin_memory"]))
num_workers=opt.workers,
pin_memory=True)
eval_dict = {'test': {'data_loader': test_data_loader}}
# model setup, model profile, optimizer config and loss definition
model = Model(backbone_type,
gd_config,
feature_dim,
num_classes=len(train_data_set.class_to_idx)).cuda()
flops, params = profile(model,
inputs=(torch.randn(1, 3, 224, 224).cuda(), ),
verbose=False)
flops, params = clever_format([flops, params])
print('# Model Params: {} FLOPs: {}'.format(params, flops))
if opt.class_loss == 'arcface':
class_criterion = losses.ArcFaceLoss(num_classes=len(train_data_set.class_to_idx), \
embedding_size=512) # , reducer=reducers.ThresholdReducer(low=0.1)
elif opt.class_loss == 'contra':
distance = distances.CosineSimilarity()
class_criterion = losses.ContrastiveLoss(distance=distance)
elif opt.class_loss == 'multi':
class_criterion = losses.MultiSimilarityLoss()
else:
class_criterion = LabelSmoothingCrossEntropyLoss(
smoothing=smoothing, temperature=temperature)
if opt.class_loss == 'arcface':
print('Using ArcFace')
optimizer = Adam([{
'params': model.parameters()
}, {
'params': class_criterion.parameters()
}],