def run_dc(epoch):
end = time.time()
model.top_layer = None
model.classifier = nn.Sequential(*list(model.classifier.children())[:-1])
# get the features for the whole dataset
features = compute_features(train_loader, model, 50000)
# cluster the features
clustering_loss = deepcluster.cluster(features)
# assign pseudo-labels
dataset_train = datasets.CIFAR10('../data',
train=True,
transform=transforms.Compose([
transforms.Pad(4),
transforms.RandomCrop(32),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
]))
train_dataset = clustering.cluster_assign(deepcluster.images_lists,
dataset_train)
# uniformly sample per target
sampler = UnifLabelSampler(int(len(train_dataset)),
deepcluster.images_lists)
train_dataloader_new = torch.utils.data.DataLoader(
train_dataset,
batch_size=args.batch_size,
num_workers=4,
sampler=sampler,
pin_memory=True,
)
# set last fully connected layer
mlp = list(model.classifier.children())
mlp.append(nn.ReLU().cuda())
model.classifier = nn.Sequential(*mlp)
# model.classifier = None
model.top_layer = nn.Linear(fd, len(deepcluster.images_lists))
model.top_layer.weight.data.normal_(0, 0.01)
model.top_layer.bias.data.zero_()
model.top_layer.cuda()
# train network with clusters as pseudo-labels
print(f'Clustering epoch {epoch} cost time: {time.time() - end} s')
loss = train(train_dataloader_new, model, criterion, optimizer, epoch)
print(f'Train Epoch: {epoch}:\tLoss: {loss}')
def main():
global args
args = parser.parse_args()
# fix random seeds
torch.manual_seed(args.seed)
torch.cuda.manual_seed_all(args.seed)
np.random.seed(args.seed)
# CNN
if args.verbose:
print('Architecture: {}'.format(args.arch))
model = models.__dict__[args.arch](sobel=args.sobel)
fd = int(model.top_layer.weight.size()[1])
model.top_layer = None
model.features = torch.nn.DataParallel(model.features)
model.cuda()
cudnn.benchmark = True
# create optimizer
optimizer = torch.optim.SGD(
filter(lambda x: x.requires_grad, model.parameters()),
lr=args.lr,
momentum=args.momentum,
weight_decay=10**args.wd,
)
# define loss function
criterion = nn.CrossEntropyLoss().cuda()
# 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']
# remove top_layer parameters from checkpoint
for key in checkpoint['state_dict']:
if 'top_layer' in key:
del checkpoint['state_dict'][key]
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))
# creating checkpoint repo
exp_check = os.path.join(args.exp, 'checkpoints')
if not os.path.isdir(exp_check):
os.makedirs(exp_check)
# creating cluster assignments log
cluster_log = Logger(os.path.join(args.exp, 'clusters'))
# preprocessing of data
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
tra = [
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(), normalize
]
# load the data
end = time.time()
dataset = datasets.ImageFolder(args.data,
transform=transforms.Compose(tra))
if args.verbose: print('Load dataset: {0:.2f} s'.format(time.time() - end))
dataloader = torch.utils.data.DataLoader(dataset,
batch_size=args.batch,
num_workers=args.workers,
pin_memory=True)
# clustering algorithm to use
deepcluster = clustering.__dict__[args.clustering](args.nmb_cluster)
# training convnet with DeepCluster
for epoch in range(args.start_epoch, args.epochs):
end = time.time()
# remove head
model.top_layer = None
model.classifier = nn.Sequential(
*list(model.classifier.children())[:-1])
# get the features for the whole dataset
features = compute_features(dataloader, model, len(dataset))
# cluster the features
clustering_loss = deepcluster.cluster(features, verbose=args.verbose)
# assign pseudo-labels
train_dataset = clustering.cluster_assign(deepcluster.images_lists,
dataset.imgs)
# uniformely sample per target
sampler = UnifLabelSampler(int(args.reassign * len(train_dataset)),
deepcluster.images_lists)
train_dataloader = torch.utils.data.DataLoader(
train_dataset,
batch_size=args.batch,
num_workers=args.workers,
sampler=sampler,
pin_memory=True,
)
# set last fully connected layer
mlp = list(model.classifier.children())
mlp.append(nn.ReLU(inplace=True).cuda())
model.classifier = nn.Sequential(*mlp)
model.top_layer = nn.Linear(fd, len(deepcluster.images_lists))
model.top_layer.weight.data.normal_(0, 0.01)
model.top_layer.bias.data.zero_()
model.top_layer.cuda()
# train network with clusters as pseudo-labels
end = time.time()
loss = train(train_dataloader, model, criterion, optimizer, epoch)
# print log
if args.verbose:
print('###### Epoch [{0}] ###### \n'
'Time: {1:.3f} s\n'
'Clustering loss: {2:.3f} \n'
'ConvNet loss: {3:.3f}'.format(epoch,
time.time() - end,
clustering_loss, loss))
try:
nmi = normalized_mutual_info_score(
clustering.arrange_clustering(deepcluster.images_lists),
clustering.arrange_clustering(cluster_log.data[-1]))
print('NMI against previous assignment: {0:.3f}'.format(nmi))
except IndexError:
pass
print('####################### \n')
# save running checkpoint
torch.save(
{
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict()
}, os.path.join(args.exp, 'checkpoint.pth.tar'))
# save cluster assignments
cluster_log.log(deepcluster.images_lists)
def main(args):
# fix random seeds
torch.manual_seed(args.seed)
torch.cuda.manual_seed_all(args.seed)
np.random.seed(args.seed)
device = torch.device('cuda:0' if torch.cuda.is_available() else "cpu")
print(device)
criterion = nn.CrossEntropyLoss()
cluster_log = Logger(os.path.join(args.exp, '../..', 'clusters.pickle'))
# CNN
if args.verbose:
print('Architecture: {}'.format(args.arch))
'''
##########################################
##########################################
# Model definition
##########################################
##########################################'''
model = models.__dict__[args.arch](bn=True,
num_cluster=args.nmb_cluster,
num_category=args.nmb_category)
fd = int(model.cluster_layer[0].weight.size()
[1]) # due to transpose, fd is input dim of W (in dim, out dim)
model.cluster_layer = None
model.category_layer = None
model.features = torch.nn.DataParallel(model.features)
model = model.double()
model.to(device)
cudnn.benchmark = True
if args.optimizer is 'Adam':
print('Adam optimizer: conv')
optimizer_body = torch.optim.Adam(
filter(lambda x: x.requires_grad, model.parameters()),
lr=args.lr_Adam,
betas=(0.9, 0.999),
weight_decay=10**args.wd,
)
else:
print('SGD optimizer: conv')
optimizer_body = torch.optim.SGD(
filter(lambda x: x.requires_grad, model.parameters()),
lr=args.lr_SGD,
momentum=args.momentum,
weight_decay=10**args.wd,
)
'''
###############
###############
category_layer
###############
###############
'''
model.category_layer = nn.Sequential(
nn.Linear(fd, args.nmb_category),
nn.Softmax(dim=1),
)
model.category_layer[0].weight.data.normal_(0, 0.01)
model.category_layer[0].bias.data.zero_()
model.category_layer = model.category_layer.double()
model.category_layer.to(device)
if args.optimizer is 'Adam':
print('Adam optimizer: conv')
optimizer_category = torch.optim.Adam(
filter(lambda x: x.requires_grad,
model.category_layer.parameters()),
lr=args.lr_Adam,
betas=(0.9, 0.999),
weight_decay=10**args.wd,
)
else:
print('SGD optimizer: conv')
optimizer_category = torch.optim.SGD(
filter(lambda x: x.requires_grad,
model.category_layer.parameters()),
lr=args.lr_SGD,
momentum=args.momentum,
weight_decay=10**args.wd,
)
'''
########################################
########################################
Create echogram sampling index
########################################
########################################'''
print('Sample echograms.')
dataset_cp, dataset_semi = sampling_echograms_full(args)
dataloader_cp = torch.utils.data.DataLoader(dataset_cp,
shuffle=False,
batch_size=args.batch,
num_workers=args.workers,
drop_last=False,
pin_memory=True)
dataloader_semi = torch.utils.data.DataLoader(dataset_semi,
shuffle=False,
batch_size=args.batch,
num_workers=args.workers,
drop_last=False,
pin_memory=True)
dataset_test = sampling_echograms_test(args)
dataloader_test = torch.utils.data.DataLoader(dataset_test,
shuffle=False,
batch_size=args.batch,
num_workers=args.workers,
drop_last=False,
pin_memory=True)
# clustering algorithm to use
deepcluster = clustering.__dict__[args.clustering](args.nmb_cluster,
args.pca)
# 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']
# remove top located layer parameters from checkpoint
copy_checkpoint_state_dict = checkpoint['state_dict'].copy()
for key in list(copy_checkpoint_state_dict):
if 'cluster_layer' in key:
del copy_checkpoint_state_dict[key]
# if 'category_layer' in key:
# del copy_checkpoint_state_dict[key]
checkpoint['state_dict'] = copy_checkpoint_state_dict
model.load_state_dict(checkpoint['state_dict'])
optimizer_body.load_state_dict(checkpoint['optimizer_body'])
optimizer_category.load_state_dict(
checkpoint['optimizer_category'])
category_save = os.path.join(args.exp, '../..',
'category_layer.pth.tar')
if os.path.isfile(category_save):
category_layer_param = torch.load(category_save)
model.category_layer.load_state_dict(category_layer_param)
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(args.resume))
# creating checkpoint repo
exp_check = os.path.join(args.exp, '../..', 'checkpoints')
if not os.path.isdir(exp_check):
os.makedirs(exp_check)
'''
#######################
#######################
PRETRAIN: commented
#######################
#######################'''
# if args.start_epoch < args.pretrain_epoch:
# if os.path.isfile(os.path.join(args.exp, '..', 'pretrain_loss_collect.pickle')):
# with open(os.path.join(args.exp, '..', 'pretrain_loss_collect.pickle'), "rb") as f:
# pretrain_loss_collect = pickle.load(f)
# else:
# pretrain_loss_collect = [[], [], [], [], []]
# print('Start pretraining with %d percent of the dataset from epoch %d/(%d)'
# % (int(args.semi_ratio * 100), args.start_epoch, args.pretrain_epoch))
# model.cluster_layer = None
#
# for epoch in range(args.start_epoch, args.pretrain_epoch):
# with torch.autograd.set_detect_anomaly(True):
# pre_loss, pre_accuracy = supervised_train(loader=dataloader_semi,
# model=model,
# crit=criterion,
# opt_body=optimizer_body,
# opt_category=optimizer_category,
# epoch=epoch, device=device, args=args)
# test_loss, test_accuracy = test(dataloader_test, model, criterion, device, args)
#
# # print log
# if args.verbose:
# print('###### Epoch [{0}] ###### \n'
# 'PRETRAIN tr_loss: {1:.3f} \n'
# 'TEST loss: {2:.3f} \n'
# 'PRETRAIN tr_accu: {3:.3f} \n'
# 'TEST accu: {4:.3f} \n'.format(epoch, pre_loss, test_loss, pre_accuracy, test_accuracy))
# pretrain_loss_collect[0].append(epoch)
# pretrain_loss_collect[1].append(pre_loss)
# pretrain_loss_collect[2].append(test_loss)
# pretrain_loss_collect[3].append(pre_accuracy)
# pretrain_loss_collect[4].append(test_accuracy)
#
# torch.save({'epoch': epoch + 1,
# 'arch': args.arch,
# 'state_dict': model.state_dict(),
# 'optimizer_body': optimizer_body.state_dict(),
# 'optimizer_category': optimizer_category.state_dict(),
# },
# os.path.join(args.exp, '..', 'checkpoint.pth.tar'))
# torch.save(model.category_layer.state_dict(), os.path.join(args.exp, '..', 'category_layer.pth.tar'))
#
# with open(os.path.join(args.exp, '..', 'pretrain_loss_collect.pickle'), "wb") as f:
# pickle.dump(pretrain_loss_collect, f)
#
# if (epoch+1) % args.checkpoints == 0:
# path = os.path.join(
# args.exp, '..',
# 'checkpoints',
# 'checkpoint_' + str(epoch) + '.pth.tar',
# )
# if args.verbose:
# print('Save checkpoint at: {0}'.format(path))
# torch.save({'epoch': epoch + 1,
# 'arch': args.arch,
# 'state_dict': model.state_dict(),
# 'optimizer_body': optimizer_body.state_dict(),
# 'optimizer_category': optimizer_category.state_dict(),
# }, path)
if os.path.isfile(os.path.join(args.exp, '../..', 'loss_collect.pickle')):
with open(os.path.join(args.exp, '../..', 'loss_collect.pickle'),
"rb") as f:
loss_collect = pickle.load(f)
else:
loss_collect = [[], [], [], [], [], [], []]
if os.path.isfile(os.path.join(args.exp, '../..', 'nmi_collect.pickle')):
with open(os.path.join(args.exp, '../..', 'nmi_collect.pickle'),
"rb") as ff:
nmi_save = pickle.load(ff)
else:
nmi_save = []
'''
#######################
#######################
MAIN TRAINING
#######################
#######################'''
for epoch in range(args.start_epoch, args.epochs):
end = time.time()
model.classifier = nn.Sequential(
*list(model.classifier.children())
[:-1]) # remove ReLU at classifier [:-1]
model.cluster_layer = None
model.category_layer = None
'''
#######################
#######################
PSEUDO-LABEL GENERATION
#######################
#######################
'''
print('Cluster the features')
features_train, input_tensors_train, labels_train = compute_features(
dataloader_cp, model, len(dataset_cp), device=device, args=args)
clustering_loss, pca_features = deepcluster.cluster(
features_train, verbose=args.verbose)
nan_location = np.isnan(pca_features)
inf_location = np.isinf(pca_features)
if (not np.allclose(nan_location, 0)) or (not np.allclose(
inf_location, 0)):
print('PCA: Feature NaN or Inf found. Nan count: ',
np.sum(nan_location), ' Inf count: ', np.sum(inf_location))
print('Skip epoch ', epoch)
torch.save(pca_features, 'tr_pca_NaN_%d.pth.tar' % epoch)
torch.save(features_train, 'tr_feature_NaN_%d.pth.tar' % epoch)
continue
print('Assign pseudo labels')
size_cluster = np.zeros(len(deepcluster.images_lists))
for i, _list in enumerate(deepcluster.images_lists):
size_cluster[i] = len(_list)
print('size in clusters: ', size_cluster)
img_label_pair_train = zip_img_label(input_tensors_train, labels_train)
train_dataset = clustering.cluster_assign(
deepcluster.images_lists,
img_label_pair_train) # Reassigned pseudolabel
# uniformly sample per target
sampler_train = UnifLabelSampler(int(len(train_dataset)),
deepcluster.images_lists)
train_dataloader = torch.utils.data.DataLoader(
train_dataset,
batch_size=args.batch,
shuffle=False,
num_workers=args.workers,
sampler=sampler_train,
pin_memory=True,
)
'''
####################################################################
####################################################################
TRSNSFORM MODEL FOR SELF-SUPERVISION // SEMI-SUPERVISION
####################################################################
####################################################################
'''
# Recover classifier with ReLU (that is not used in clustering)
mlp = list(model.classifier.children(
)) # classifier that ends with linear(512 * 128). No ReLU at the end
mlp.append(nn.ReLU(inplace=True).to(device))
model.classifier = nn.Sequential(*mlp)
model.classifier.to(device)
'''SELF-SUPERVISION (PSEUDO-LABELS)'''
model.category_layer = None
model.cluster_layer = nn.Sequential(
nn.Linear(fd, args.nmb_cluster), # nn.Linear(4096, num_cluster),
nn.Softmax(
dim=1
), # should be removed and replaced by ReLU for category_layer
)
model.cluster_layer[0].weight.data.normal_(0, 0.01)
model.cluster_layer[0].bias.data.zero_()
model.cluster_layer = model.cluster_layer.double()
model.cluster_layer.to(device)
''' train network with clusters as pseudo-labels '''
with torch.autograd.set_detect_anomaly(True):
pseudo_loss, semi_loss, semi_accuracy = semi_train(
train_dataloader,
dataloader_semi,
model,
fd,
criterion,
optimizer_body,
optimizer_category,
epoch,
device=device,
args=args)
# save checkpoint
if (epoch + 1) % args.checkpoints == 0:
path = os.path.join(
args.exp,
'../..',
'checkpoints',
'checkpoint_' + str(epoch) + '.pth.tar',
)
if args.verbose:
print('Save checkpoint at: {0}'.format(path))
torch.save(
{
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': model.state_dict(),
'optimizer_body': optimizer_body.state_dict(),
'optimizer_category': optimizer_category.state_dict(),
}, path)
'''
##############
##############
# TEST phase
##############
##############
'''
test_loss, test_accuracy, test_pred, test_label = test(
dataloader_test, model, criterion, device, args)
'''Save prediction of the test set'''
if (epoch % args.save_epoch == 0):
with open(
os.path.join(args.exp, '../..',
'sup_epoch_%d_te.pickle' % epoch), "wb") as f:
pickle.dump([test_pred, test_label], f)
if args.verbose:
print('###### Epoch [{0}] ###### \n'
'Time: {1:.3f} s\n'
'Pseudo tr_loss: {2:.3f} \n'
'SEMI tr_loss: {3:.3f} \n'
'TEST loss: {4:.3f} \n'
'Clustering loss: {5:.3f} \n'
'SEMI accu: {6:.3f} \n'
'TEST accu: {7:.3f} \n'.format(epoch,
time.time() - end,
pseudo_loss, semi_loss,
test_loss, clustering_loss,
semi_accuracy, test_accuracy))
try:
nmi = normalized_mutual_info_score(
clustering.arrange_clustering(deepcluster.images_lists),
clustering.arrange_clustering(cluster_log.data[-1]))
nmi_save.append(nmi)
print('NMI against previous assignment: {0:.3f}'.format(nmi))
with open(
os.path.join(args.exp, '../..', 'nmi_collect.pickle'),
"wb") as ff:
pickle.dump(nmi_save, ff)
except IndexError:
pass
print('####################### \n')
# save cluster assignments
cluster_log.log(deepcluster.images_lists)
# save running checkpoint
torch.save(
{
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': model.state_dict(),
'optimizer_body': optimizer_body.state_dict(),
'optimizer_category': optimizer_category.state_dict(),
}, os.path.join(args.exp, '../..', 'checkpoint.pth.tar'))
torch.save(model.category_layer.state_dict(),
os.path.join(args.exp, '../..', 'category_layer.pth.tar'))
loss_collect[0].append(epoch)
loss_collect[1].append(pseudo_loss)
loss_collect[2].append(semi_loss)
loss_collect[3].append(clustering_loss)
loss_collect[4].append(test_loss)
loss_collect[5].append(semi_accuracy)
loss_collect[6].append(test_accuracy)
with open(os.path.join(args.exp, '../..', 'loss_collect.pickle'),
"wb") as f:
pickle.dump(loss_collect, f)
'''
############################
############################
# PSEUDO-LABEL GEN: Test set
############################
############################
'''
model.classifier = nn.Sequential(
*list(model.classifier.children())
[:-1]) # remove ReLU at classifier [:-1]
model.cluster_layer = None
model.category_layer = None
print('TEST set: Cluster the features')
features_te, input_tensors_te, labels_te = compute_features(
dataloader_test,
model,
len(dataset_test),
device=device,
args=args)
clustering_loss_te, pca_features_te = deepcluster.cluster(
features_te, verbose=args.verbose)
mlp = list(model.classifier.children(
)) # classifier that ends with linear(512 * 128). No ReLU at the end
mlp.append(nn.ReLU(inplace=True).to(device))
model.classifier = nn.Sequential(*mlp)
model.classifier.to(device)
nan_location = np.isnan(pca_features_te)
inf_location = np.isinf(pca_features_te)
if (not np.allclose(nan_location, 0)) or (not np.allclose(
inf_location, 0)):
print('PCA: Feature NaN or Inf found. Nan count: ',
np.sum(nan_location), ' Inf count: ', np.sum(inf_location))
print('Skip epoch ', epoch)
torch.save(pca_features_te, 'te_pca_NaN_%d.pth.tar' % epoch)
torch.save(features_te, 'te_feature_NaN_%d.pth.tar' % epoch)
continue
# save patches per epochs
cp_epoch_out = [
features_te, deepcluster.images_lists,
deepcluster.images_dist_lists, input_tensors_te, labels_te
]
if (epoch % args.save_epoch == 0):
with open(
os.path.join(args.exp, '../..',
'cp_epoch_%d_te.pickle' % epoch), "wb") as f:
pickle.dump(cp_epoch_out, f)
with open(
os.path.join(args.exp, '../..',
'pca_epoch_%d_te.pickle' % epoch), "wb") as f:
pickle.dump(pca_features_te, f)
def main(args):
# fix random seeds
torch.manual_seed(args.seed)
torch.cuda.manual_seed_all(args.seed)
np.random.seed(args.seed)
device = torch.device('cuda:0' if torch.cuda.is_available() else "cpu")
print(device)
# CNN
if args.verbose:
print('Architecture: {}'.format(args.arch))
model = models.__dict__[args.arch](sobel=False,
bn=True,
out=args.nmb_cluster)
fd = int(model.top_layer[0].weight.size()
[1]) # due to transpose, fd is input dim of W (in dim, out dim)
model.top_layer = None
model.features = torch.nn.DataParallel(model.features)
model = model.double()
model.to(device)
cudnn.benchmark = True
# create optimizer
# optimizer = torch.optim.SGD(
# filter(lambda x: x.requires_grad, model.parameters()),
# lr=args.lr,
# momentum=args.momentum,
# weight_decay=10**args.wd,
# )
optimizer = torch.optim.Adam(
filter(lambda x: x.requires_grad, model.parameters()),
lr=args.lr,
betas=(0.5, 0.99),
weight_decay=10**args.wd,
)
criterion = nn.CrossEntropyLoss()
# 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']
# remove top_layer parameters from checkpoint
copy_checkpoint_state_dict = checkpoint['state_dict'].copy()
for key in list(copy_checkpoint_state_dict):
if 'top_layer' in key:
del copy_checkpoint_state_dict[key]
checkpoint['state_dict'] = copy_checkpoint_state_dict
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))
# creating checkpoint repo
exp_check = os.path.join(args.exp, '../checkpoints')
if not os.path.isdir(exp_check):
os.makedirs(exp_check)
# creating cluster assignments log
cluster_log = Logger(os.path.join(args.exp, 'clusters.pickle'))
# load dataset (initial echograms)
window_size = [args.window_dim, args.window_dim]
# # Create echogram sampling index
print('Sample echograms.')
end = time.time()
dataset_cp = sampling_echograms_full(args)
dataloader_cp = torch.utils.data.DataLoader(dataset_cp,
shuffle=False,
batch_size=args.batch,
num_workers=args.workers,
drop_last=False,
pin_memory=True)
if args.verbose:
print('Load dataset: {0:.2f} s'.format(time.time() - end))
# clustering algorithm to use
deepcluster = clustering.__dict__[args.clustering](args.nmb_cluster,
args.pca)
# deepcluster = clustering.Kmeans(no.cluster, dim.pca)
loss_collect = [[], [], []]
# training convnet with DeepCluster
for epoch in range(args.start_epoch, args.epochs):
# remove head
model.top_layer = None
model.classifier = nn.Sequential(*list(model.classifier.children(
))) # End with linear(512*128) in original vgg)
# ReLU in .classfier() will follow later
# get the features for the whole dataset
features_train, input_tensors_train, labels_train = compute_features(
dataloader_cp, model, len(dataset_cp), device=device, args=args)
# cluster the features
print('Cluster the features')
end = time.time()
clustering_loss = deepcluster.cluster(features_train,
verbose=args.verbose)
print('Cluster time: {0:.2f} s'.format(time.time() - end))
# save patches per epochs
if ((epoch + 1) % args.save_epoch == 0):
end = time.time()
cp_epoch_out = [
features_train, deepcluster.images_lists,
deepcluster.images_dist_lists, input_tensors_train,
labels_train
]
with open("./cp_epoch_%d.pickle" % epoch, "wb") as f:
pickle.dump(cp_epoch_out, f)
print('Feature save time: {0:.2f} s'.format(time.time() - end))
# assign pseudo-labels
print('Assign pseudo labels')
size_cluster = np.zeros(len(deepcluster.images_lists))
for i, _list in enumerate(deepcluster.images_lists):
size_cluster[i] = len(_list)
print('size in clusters: ', size_cluster)
img_label_pair_train = zip_img_label(input_tensors_train, labels_train)
train_dataset = clustering.cluster_assign(
deepcluster.images_lists,
img_label_pair_train) # Reassigned pseudolabel
# ((img[imgidx], label[imgidx]), pseudolabel, imgidx)
# N = len(imgidx)
# uniformly sample per target
sampler_train = UnifLabelSampler(int(len(train_dataset)),
deepcluster.images_lists)
train_dataloader = torch.utils.data.DataLoader(
train_dataset,
batch_size=args.batch,
shuffle=False,
num_workers=args.workers,
sampler=sampler_train,
pin_memory=True,
)
# set last fully connected layer
mlp = list(model.classifier.children()
) # classifier that ends with linear(512 * 128)
mlp.append(nn.ReLU().to(device))
model.classifier = nn.Sequential(*mlp)
model.top_layer = nn.Sequential(
nn.Linear(fd, args.nmb_cluster),
nn.Softmax(dim=1),
)
# model.top_layer = nn.Linear(fd, args.nmb_cluster)
model.top_layer[0].weight.data.normal_(0, 0.01)
model.top_layer[0].bias.data.zero_()
model.top_layer = model.top_layer.double()
model.top_layer.to(device)
# train network with clusters as pseudo-labels
end = time.time()
with torch.autograd.set_detect_anomaly(True):
loss, tr_epoch_out = train(train_dataloader,
model,
criterion,
optimizer,
epoch,
device=device,
args=args)
print('Train time: {0:.2f} s'.format(time.time() - end))
if ((epoch + 1) % args.save_epoch == 0):
end = time.time()
with open("./tr_epoch_%d.pickle" % epoch, "wb") as f:
pickle.dump(tr_epoch_out, f)
print('Save train time: {0:.2f} s'.format(time.time() - end))
# Accuracy with training set (output vs. pseudo label)
accuracy_tr = np.mean(
tr_epoch_out[1] == np.argmax(tr_epoch_out[2], axis=1))
# print log
if args.verbose:
print('###### Epoch [{0}] ###### \n'
'Time: {1:.3f} s\n'
'Clustering loss: {2:.3f} \n'
'ConvNet tr_loss: {3:.3f} \n'
'ConvNet tr_acc: {4:.3f} \n'.format(epoch,
time.time() - end,
clustering_loss, loss,
accuracy_tr))
try:
nmi = normalized_mutual_info_score(
clustering.arrange_clustering(deepcluster.images_lists),
clustering.arrange_clustering(cluster_log.data[-1]))
print('NMI against previous assignment: {0:.3f}'.format(nmi))
except IndexError:
pass
print('####################### \n')
# save running checkpoint
torch.save(
{
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict()
}, os.path.join(args.exp, 'checkpoint.pth.tar'))
loss_collect[0].append(epoch)
loss_collect[1].append(loss)
loss_collect[2].append(accuracy_tr)
with open("./loss_collect.pickle", "wb") as f:
pickle.dump(loss_collect, f)
# save cluster assignments
cluster_log.log(deepcluster.images_lists)
def main(args):
# fix random seeds
torch.manual_seed(args.seed)
torch.cuda.manual_seed_all(args.seed)
np.random.seed(args.seed)
# CNN
if args.verbose:
print('Architecture: {}'.format(args.arch))
model = models.__dict__[args.arch](sobel=args.sobel)
fd = int(model.top_layer.weight.size()[1])
model.top_layer = None
model.features = torch.nn.DataParallel(model.features)
model.cuda()
cudnn.benchmark = True
print('CNN builded.')
# create optimizer
optimizer = torch.optim.SGD(
filter(lambda x: x.requires_grad, model.parameters()),
lr=args.lr,
momentum=args.momentum,
weight_decay=10**args.wd,
)
print('Optimizer created.')
# define loss function
criterion = nn.CrossEntropyLoss().cuda()
# 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']
# remove top_layer parameters from checkpoint
for key in list(checkpoint['state_dict']):
if 'top_layer' in key:
del checkpoint['state_dict'][key]
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))
# creating checkpoint repo
exp_check = os.path.join(args.exp, 'checkpoints')
if not os.path.isdir(exp_check):
os.makedirs(exp_check)
# creating cluster assignments log
cluster_log = Logger(os.path.join(args.exp, 'clusters'))
epochs_log = Logger(os.path.join(args.exp, 'epochs300'))
# Loading and preprocessing of data: custom Rescale and ToTensor transformations for VidDataset.
# VidDataset has a box_frame, which is a pandas Dataframe containing images path an their bb coordinates.
# Each VidDataset sample is a dict formed by a tensor (the image) and crop_coord (bb xmin, xmax, ymin, ymax).
# If a pickled dataset is passed, it will be deserialized and used, else it will be normally loaded.
# It is useful when we want to preprocess a dataset.
print('Start loading dataset...')
end = time.time()
if args.dataset_pkl:
dataset = deserialize_obj(args.dataset_pkl)
# I will never use labels in deepcluster
dataset.vid_labels = None
else:
tra = [preprocessing.Rescale((224, 224)), preprocessing.ToTensor()]
dataset = VidDataset(xml_annotations_dir=args.ann,
root_dir=args.data,
transform=transforms.Compose(tra))
dataset.imgs = dataset.imgs[0::args.load_step]
dataset.samples = dataset.samples[0::args.load_step]
print('Load dataset: {0:.2f} s'.format(time.time() - end))
dataloader = torch.utils.data.DataLoader(dataset,
batch_size=args.batch,
num_workers=args.workers,
pin_memory=True)
# calculate batch size sum (better clean-up data with clean_data.py)
dataset_len = 0
if not args.dataset_pkl:
dataloader.collate_fn = my_collate
for s in dataloader:
dataset_len += len(s['image'])
else:
dataset_len = len(dataset.imgs)
print("Dataset final dimension: ", dataset_len)
# clustering algorithm to use
deepcluster = clustering.__dict__[args.clustering](args.nmb_cluster)
# training convnet with DeepCluster
for epoch in range(args.start_epoch, args.epochs):
end = time.time()
# remove head
model.top_layer = None
model.classifier = nn.Sequential(
*list(model.classifier.children())[:-1])
# get the features for the whole dataset hardcoded dataset dim for step=5
features = compute_features(dataloader, model, args.load_step,
dataset_len)
# cluster the features
if args.verbose:
print('Cluster the features')
clustering_loss = deepcluster.cluster(features, verbose=args.verbose)
# assign pseudo-labels
if args.verbose:
print('Assign pseudo labels')
train_dataset = clustering.cluster_assign(deepcluster.images_lists,
dataset.imgs)
# uniformly sample per target
sampler = UnifLabelSampler(int(args.reassign * len(train_dataset)),
deepcluster.images_lists)
train_dataloader = torch.utils.data.DataLoader(
train_dataset,
batch_size=args.batch,
num_workers=args.workers,
sampler=sampler,
pin_memory=True,
)
if not args.dataset_pkl:
train_dataloader.collate_fn = my_collate
# set last fully connected layer
mlp = list(model.classifier.children())
mlp.append(nn.ReLU(inplace=True).cuda())
model.classifier = nn.Sequential(*mlp)
model.top_layer = nn.Linear(fd, len(deepcluster.images_lists))
model.top_layer.weight.data.normal_(0, 0.01)
model.top_layer.bias.data.zero_()
model.top_layer.cuda()
# train network with clusters as pseudo-labels
end = time.time()
loss = train(train_dataloader, model, criterion, optimizer, epoch)
# print log
if args.verbose:
print('###### Epoch [{0}] ###### \n'
'Time: {1:.3f} s\n'
'Clustering loss: {2:.3f} \n'
'ConvNet loss: {3:.3f}'.format(epoch,
time.time() - end,
clustering_loss, loss))
epoch_log = [epoch, time.time() - end, clustering_loss, loss]
try:
nmi = normalized_mutual_info_score(
clustering.arrange_clustering(deepcluster.images_lists),
clustering.arrange_clustering(cluster_log.data[-1]))
print('NMI against previous assignment: {0:.3f}'.format(nmi))
epoch_log.append(nmi)
except IndexError:
pass
print('####################### \n')
# save running checkpoint
torch.save(
{
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict()
}, os.path.join(args.exp, 'checkpoint.pth.tar'))
# save cluster assignments
cluster_log.log(deepcluster.images_lists)
epochs_log.log(epoch_log)
def main():
global args
args = parser.parse_args()
# fix random seeds
torch.manual_seed(args.seed)
torch.cuda.manual_seed_all(args.seed)
np.random.seed(args.seed)
logs = []
# CNN
if args.verbose:
print('Architecture: {}'.format(args.arch))
if args.arch == 'inceptionv1':
model = models.__dict__[args.arch](
sobel=args.sobel,
weight_file='/home/farbod/honours/convert/kit_pytorch.npy')
else:
model = models.__dict__[args.arch](sobel=args.sobel)
fd = int(model.top_layer.weight.size()[1])
model.top_layer = None
if args.arch == 'inceptionv1':
for key in model.modules():
if isinstance(key, nn.Module): continue
key = torch.nn.DataParallel(key).cuda()
else:
model.features = torch.nn.DataParallel(model.features)
model.cuda()
cudnn.benchmark = True
#for param in model.parameters():
# param.requires_grad = False
#for param in model.classifier.parameters():
# param.requires_grad = True
# create optimizer
optimizer = torch.optim.SGD(
filter(lambda x: x.requires_grad, model.parameters()),
lr=args.lr,
momentum=args.momentum,
weight_decay=10**args.wd,
)
# define loss function
criterion = nn.CrossEntropyLoss().cuda()
# 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']
# remove top_layer parameters from checkpoint
for key in checkpoint['state_dict']:
if 'top_layer' in key:
del checkpoint['state_dict'][key]
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))
# creating checkpoint repo
exp_check = os.path.join(args.exp, 'checkpoints')
if not os.path.isdir(exp_check):
os.makedirs(exp_check)
plot_dir = os.path.join(args.exp, 'plots')
if not os.path.isdir(plot_dir):
os.makedirs(plot_dir)
# creating cluster assignments log
cluster_log = Logger(os.path.join(args.exp, 'clusters'))
# preprocessing of data
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
#normalize = transforms.Normalize(mean=[0.5, 0.5, 0.5],
# std=[0.5, 0.5, 0.5])
tra = [
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(), normalize
]
# load the data
end = time.time()
dataset = datasets.ImageFolder(args.data,
transform=transforms.Compose(tra))
if args.verbose: print('Load dataset: {0:.2f} s'.format(time.time() - end))
dataloader = torch.utils.data.DataLoader(dataset,
batch_size=args.batch,
num_workers=args.workers,
pin_memory=True)
# get ground truth labels for nmi
num_classes = 65
labels = [[] for i in range(num_classes)]
for i, (_, label) in enumerate(dataset.imgs):
labels[label].append(i)
# clustering algorithm to use
deepcluster = clustering.__dict__[args.clustering](args.nmb_cluster)
# training convnet with DeepCluster
for epoch in range(args.start_epoch, args.epochs):
end = time.time()
# remove head
model.top_layer = None
model.classifier = nn.Sequential(
*list(model.classifier.children())[:-1])
# get the features for the whole dataset
features = compute_features(dataloader, model, len(dataset))
# cluster the features
clustering_loss, plot, davg = deepcluster.cluster(features,
verbose=args.verbose)
print davg
if epoch < 20:
plot.savefig(os.path.join(plot_dir, 'e{}'.format(epoch)))
# assign pseudo-labels
train_dataset = clustering.cluster_assign(deepcluster.images_lists,
dataset.imgs)
#for i, image in enumerate(train_dataset):
# save_dir = os.path.join('./viz_emb_start', str(image[1]))
# if not os.path.isdir(save_dir):
# os.makedirs(save_dir)
# imn = (image[0].data.cpu().numpy() * 112) + 112
# imn = np.swapaxes(imn, 0, 2)
# imn = np.swapaxes(imn, 1, 0)
# #print imn.astype('uint8')
# #print imn.astype('uint8').shape
# im = Image.fromarray(imn.astype('uint8'))
# im.save(os.path.join(save_dir, '{}.jpg'.format(i)))
# uniformely sample per target
sampler = UnifLabelSampler(int(args.reassign * len(train_dataset)),
deepcluster.images_lists)
train_dataloader = torch.utils.data.DataLoader(
train_dataset,
batch_size=args.batch,
num_workers=args.workers,
sampler=sampler,
pin_memory=True,
)
# set last fully connected layer
mlp = list(model.classifier.children())
mlp.append(nn.ReLU(inplace=True).cuda())
model.classifier = nn.Sequential(*mlp)
model.top_layer = nn.Linear(fd, len(deepcluster.images_lists))
model.top_layer.weight.data.normal_(0, 0.01)
model.top_layer.bias.data.zero_()
model.top_layer.cuda()
# train network with clusters as pseudo-labels
end = time.time()
loss = train(train_dataloader, model, criterion, optimizer, epoch)
# print log
if args.verbose:
print(
'###### Epoch [{0}] ###### \n'
'Time: {1:.3f} s\n'
'Clustering loss: {2:.3f} \n'
'ConvNet loss: {3:.3f}'.format(epoch,
time.time() - end,
clustering_loss, loss))
nmi_prev = 0
nmi_gt = 0
try:
nmi_prev = normalized_mutual_info_score(
clustering.arrange_clustering(deepcluster.images_lists),
clustering.arrange_clustering(cluster_log.data[-1]))
print('NMI against previous assignment: {0:.3f}'.format(
nmi_prev))
except IndexError:
pass
nmi_gt = normalized_mutual_info_score(
clustering.arrange_clustering(deepcluster.images_lists),
clustering.arrange_clustering(labels))
print('NMI against ground-truth labels: {0:.3f}'.format(nmi_gt))
print('####################### \n')
logs.append([epoch, clustering_loss, loss, nmi_prev, nmi_gt, davg])
# save running checkpoint
if (epoch + 1) % 10 == 0 or epoch == 0:
torch.save(
{
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict()
},
os.path.join(args.exp,
'checkpoint_{}.pth.tar'.format(epoch + 1)))
# save cluster assignments
cluster_log.log(deepcluster.images_lists)
scipy.io.savemat(os.path.join(args.exp, 'logs.mat'),
{'logs': np.array(logs)})
def main(args):
# fix random seeds
torch.manual_seed(args.seed)
torch.cuda.manual_seed_all(args.seed)
np.random.seed(args.seed)
device = torch.device('cuda:0' if torch.cuda.is_available() else "cpu")
print(device)
criterion_pseudo = nn.CrossEntropyLoss()
criterion_sup = nn.CrossEntropyLoss(ignore_index=-1,
weight=torch.Tensor([10, 300, 250]).to(
device=device, dtype=torch.double))
# CNN
if args.verbose:
print('Architecture: {}'.format(args.arch))
'''
##########################################
##########################################
# Model definition
##########################################
##########################################'''
model = models.__dict__[args.arch](bn=True,
num_cluster=args.nmb_cluster,
num_category=args.nmb_category)
fd = int(model.cluster_layer[0].weight.size()
[1]) # due to transpose, fd is input dim of W (in dim, out dim)
model.cluster_layer = None
model.category_layer = None
model.features = torch.nn.DataParallel(model.features)
model.to(device, dtype=torch.double)
cudnn.benchmark = True
if args.optimizer is 'Adam':
print('Adam optimizer: conv')
optimizer_body = torch.optim.Adam(
filter(lambda x: x.requires_grad, model.parameters()),
lr=args.lr_Adam,
betas=(0.9, 0.999),
weight_decay=10**args.wd,
)
else:
print('SGD optimizer: conv')
optimizer_body = torch.optim.SGD(
filter(lambda x: x.requires_grad, model.parameters()),
lr=args.lr_SGD,
momentum=args.momentum,
weight_decay=10**args.wd,
)
'''
###############
###############
category_layer
###############
###############
'''
model.category_layer = nn.Sequential(
nn.Linear(fd, args.nmb_category),
nn.Softmax(dim=1),
)
model.category_layer[0].weight.data.normal_(0, 0.01)
model.category_layer[0].bias.data.zero_()
model.category_layer.to(device, dtype=torch.double)
if args.optimizer is 'Adam':
print('Adam optimizer: conv')
optimizer_category = torch.optim.Adam(
filter(lambda x: x.requires_grad,
model.category_layer.parameters()),
lr=args.lr_Adam,
betas=(0.9, 0.999),
weight_decay=10**args.wd,
)
else:
print('SGD optimizer: conv')
optimizer_category = torch.optim.SGD(
filter(lambda x: x.requires_grad,
model.category_layer.parameters()),
lr=args.lr_SGD,
momentum=args.momentum,
weight_decay=10**args.wd,
)
'''
########################################
########################################
Create echogram sampling index
########################################
########################################'''
print('Sample echograms.')
dataset_cp, dataset_semi = sampling_echograms_full_for_comparisonP2(
args) # For comparison (paper #2)
dataloader_cp = torch.utils.data.DataLoader(dataset_cp,
shuffle=False,
batch_size=args.batch,
num_workers=args.workers,
drop_last=False,
pin_memory=True)
dataloader_semi = torch.utils.data.DataLoader(dataset_semi,
shuffle=False,
batch_size=args.batch,
num_workers=args.workers,
drop_last=False,
pin_memory=True)
dataset_te = sampling_echograms_test_for_comparisonP2()
dataloader_test = torch.utils.data.DataLoader(dataset_te,
shuffle=False,
batch_size=args.batch,
num_workers=args.workers,
drop_last=False,
pin_memory=True)
dataset_2019, label_2019, patch_loc = sampling_echograms_2019_for_comparisonP2(
)
dataloader_2019 = torch.utils.data.DataLoader(
dataset_2019,
batch_size=1,
shuffle=False,
num_workers=args.workers,
worker_init_fn=np.random.seed,
drop_last=False,
pin_memory=True)
deepcluster = clustering.__dict__[args.clustering](args.nmb_cluster,
args.pca)
# 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']
# remove top located layer parameters from checkpoint
copy_checkpoint_state_dict = checkpoint['state_dict'].copy()
for key in list(copy_checkpoint_state_dict):
if 'cluster_layer' in key:
del copy_checkpoint_state_dict[key]
checkpoint['state_dict'] = copy_checkpoint_state_dict
model.load_state_dict(checkpoint['state_dict'])
optimizer_body.load_state_dict(checkpoint['optimizer_body'])
optimizer_category.load_state_dict(
checkpoint['optimizer_category'])
category_save = os.path.join(args.exp, 'category_layer.pth.tar')
if os.path.isfile(category_save):
category_layer_param = torch.load(category_save)
model.category_layer.load_state_dict(category_layer_param)
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(args.resume))
# creating checkpoint repo
exp_check = os.path.join(args.exp, 'checkpoints')
if not os.path.isdir(exp_check):
os.makedirs(exp_check)
exp_test = os.path.join(args.exp, 'test')
for dir_2 in ['2019', 'pred']:
dir_to_make = os.path.join(exp_test, dir_2)
if not os.path.isdir(dir_to_make):
os.makedirs(dir_to_make)
'''
#######################
#######################
MAIN TRAINING
#######################
#######################'''
for epoch in range(args.start_epoch, args.epochs):
print(
'##################### Start training at Epoch %d ################'
% epoch)
model.classifier = nn.Sequential(
*list(model.classifier.children())
[:-1]) # remove ReLU at classifier [:-1]
model.cluster_layer = None
model.category_layer = None
'''
#######################
#######################
PSEUDO-LABEL GENERATION
#######################
#######################
'''
print('Cluster the features')
features_train, input_tensors_train, labels_train = compute_features_for_comparisonP2(
dataloader_cp,
model,
len(dataset_cp) * args.for_comparisonP2_batchsize,
device=device,
args=args)
clustering_loss, pca_features = deepcluster.cluster(
features_train, verbose=args.verbose)
nan_location = np.isnan(pca_features)
inf_location = np.isinf(pca_features)
if (not np.allclose(nan_location, 0)) or (not np.allclose(
inf_location, 0)):
print('PCA: Feature NaN or Inf found. Nan count: ',
np.sum(nan_location), ' Inf count: ', np.sum(inf_location))
print('Skip epoch ', epoch)
torch.save(pca_features, 'tr_pca_NaN_%d.pth.tar' % epoch)
torch.save(features_train, 'tr_feature_NaN_%d.pth.tar' % epoch)
continue
print('Assign pseudo labels')
size_cluster = np.zeros(len(deepcluster.images_lists))
for i, _list in enumerate(deepcluster.images_lists):
size_cluster[i] = len(_list)
print('size in clusters: ', size_cluster)
img_label_pair_train = zip_img_label(input_tensors_train, labels_train)
train_dataset = clustering.cluster_assign(
deepcluster.images_lists,
img_label_pair_train) # Reassigned pseudolabel
# uniformly sample per target
sampler_train = UnifLabelSampler(int(len(train_dataset)),
deepcluster.images_lists)
train_dataloader = torch.utils.data.DataLoader(
train_dataset,
batch_size=args.for_comparisonP2_batchsize, #args.batch
shuffle=False,
num_workers=args.workers,
sampler=sampler_train,
pin_memory=True,
)
'''
####################################################################
####################################################################
TRSNSFORM MODEL FOR SELF-SUPERVISION // SEMI-SUPERVISION
####################################################################
####################################################################
'''
# Recover classifier with ReLU (that is not used in clustering)
mlp = list(model.classifier.children(
)) # classifier that ends with linear(512 * 128). No ReLU at the end
mlp.append(nn.ReLU(inplace=True).to(device))
model.classifier = nn.Sequential(*mlp)
model.classifier.to(device=device, dtype=torch.double)
'''SELF-SUPERVISION (PSEUDO-LABELS)'''
model.category_layer = None
model.cluster_layer = nn.Sequential(
nn.Linear(fd, args.nmb_cluster), # nn.Linear(4096, num_cluster),
nn.Softmax(
dim=1
), # should be removed and replaced by ReLU for category_layer
)
model.cluster_layer[0].weight.data.normal_(0, 0.01)
model.cluster_layer[0].bias.data.zero_()
# model.cluster_layer = model.cluster_layer.double()
model.cluster_layer.to(device=device, dtype=torch.double)
''' train network with clusters as pseudo-labels '''
with torch.autograd.set_detect_anomaly(True):
pseudo_loss, semi_loss, semi_accuracy = semi_train_for_comparisonP2(
train_dataloader,
dataloader_semi,
model,
fd,
criterion_pseudo,
criterion_sup,
optimizer_body,
optimizer_category,
epoch,
device=device,
args=args)
# save checkpoint
if epoch % args.checkpoints == 0:
path = os.path.join(
args.exp,
'checkpoints',
str(epoch) + '_checkpoint.pth.tar',
)
if args.verbose:
print('Save checkpoint at: {0}'.format(path))
torch.save(
{
'epoch': epoch,
'arch': args.arch,
'state_dict': model.state_dict(),
'optimizer_body': optimizer_body.state_dict(),
'optimizer_category': optimizer_category.state_dict(),
}, path)
torch.save(
model.category_layer.state_dict(),
os.path.join(args.exp, 'checkpoints',
'%d_category_layer.pth.tar' % epoch))
# save running checkpoint
torch.save(
{
'epoch': epoch,
'arch': args.arch,
'state_dict': model.state_dict(),
'optimizer_body': optimizer_body.state_dict(),
'optimizer_category': optimizer_category.state_dict(),
}, os.path.join(args.exp, 'checkpoint.pth.tar'))
torch.save(model.category_layer.state_dict(),
os.path.join(args.exp, 'category_layer.pth.tar'))
'''
##############
##############
# TEST phase
##############
##############
'''
test_loss, test_accuracy, test_pred, test_label, test_pred_softmax = test_for_comparisonP2(
dataloader_test, model, criterion_sup, device, args)
test_pred_large = rebuild_pred_patch(test_pred)
test_softmax_large = rebuild_pred_patch(test_pred_softmax)
test_label_large = rebuild_pred_patch(test_label)
'''Save prediction of the test set'''
if (epoch % args.save_epoch == 0):
with open(
os.path.join(
args.exp, 'test', 'pred',
'pred_softmax_label_epoch_%d_te.pickle' % epoch),
"wb") as f:
pickle.dump(
[test_pred_large, test_softmax_large, test_label_large], f)
fpr, \
tpr, \
roc_auc, \
roc_auc_macro, \
prob_mat, \
mat, \
f1_score, \
kappa, \
bg_accu, \
se_accu, \
ot_accu = test_analysis(test_pred_large, test_softmax_large, epoch, args)
if os.path.isfile(
os.path.join(args.exp, 'records_te_epoch_patch.pth.tar')):
records_te_epoch = torch.load(
os.path.join(args.exp, 'records_te_epoch_patch.pth.tar'))
else:
records_te_epoch = {
'epoch': [],
'fpr': [],
'tpr': [],
'roc_auc': [],
'roc_auc_macro': [],
'prob_mat': [],
'mat': [],
'f1_score': [],
'kappa': [],
'BG_accu_epoch': [],
'SE_accu_epoch': [],
'OT_accu_epoch': [],
}
records_te_epoch['epoch'].append(epoch)
records_te_epoch['fpr'].append(fpr)
records_te_epoch['tpr'].append(tpr)
records_te_epoch['roc_auc'].append(roc_auc)
records_te_epoch['roc_auc_macro'].append(roc_auc_macro)
records_te_epoch['prob_mat'].append(prob_mat)
records_te_epoch['mat'].append(mat)
records_te_epoch['f1_score'].append(f1_score)
records_te_epoch['kappa'].append(kappa)
records_te_epoch['BG_accu_epoch'].append(bg_accu)
records_te_epoch['SE_accu_epoch'].append(se_accu)
records_te_epoch['OT_accu_epoch'].append(ot_accu)
torch.save(records_te_epoch,
os.path.join(args.exp, 'records_te_epoch_patch.pth.tar'))
'''
##############
##############
# 2019 phase
##############
##############
'''
test_loss_2019, test_accuracy_2019, test_pred_2019, test_label_2019, test_pred_softmax_2019 = test_for_comparisonP2(
dataloader_2019, model, criterion_sup, device, args)
test_pred_large_2019 = rebuild_pred_patch(test_pred_2019)
test_softmax_large_2019 = rebuild_pred_patch(test_pred_softmax_2019)
test_label_large_2019 = rebuild_pred_patch(test_label_2019)
test_and_plot_2019(test_pred_large_2019, test_label_large_2019, epoch,
args)
def main():
global args
use_cuda = torch.cuda.is_available()
device = torch.device("cuda" if use_cuda else "cpu")
criterion = nn.CrossEntropyLoss()
# fix random seeds
torch.manual_seed(args.seed)
torch.cuda.manual_seed_all(args.seed)
np.random.seed(args.seed)
# CNN
if args.verbose:
print('Architecture: VGGMiniCBR')
model = VGGMiniCBR(num_classes=10)
fd = int(model.top_layer.weight.size()[1])
model.top_layer = None
model.to(device)
cudnn.benchmark = True
# create optimizer
optimizer = torch.optim.SGD(
filter(lambda x: x.requires_grad, model.parameters()),
lr=args.lr,
momentum=args.momentum,
weight_decay=10 ** args.wd,
)
# optimizer = torch.optim.Adam(filter(lambda x: x.requires_grad, model.parameters()))
# creating checkpoint repo
exp_check = os.path.join(args.exp, 'checkpoints')
if not os.path.isdir(exp_check):
os.makedirs(exp_check)
cluster_log = Logger(os.path.join(exp_path, 'clusters'))
tra = [
transforms.Grayscale(num_output_channels=1),
transforms.RandomAffine(degrees=5, translate=(0.03, 0.03), scale=(0.95, 1.05), shear=5),
transforms.ToTensor(),
transforms.Normalize((mean_std[use_zca][0],), (mean_std[use_zca][1],))
]
end = time.time()
dataset = datasets.ImageFolder(args.data, transform=transforms.Compose(tra))
if args.verbose:
print('Load dataset: {0:.2f} s'.format(time.time() - end))
dataloader = torch.utils.data.DataLoader(dataset,
batch_size=args.batch,
num_workers=args.workers,
pin_memory=True)
# clustering algorithm to use
deepcluster = clustering.__dict__[args.clustering](args.nmb_cluster)
# training convnet with DeepCluster
for epoch in range(args.start_epoch, args.epochs):
end = time.time()
# remove head
model.top_layer = None
model.classifier = nn.Sequential(*list(model.classifier.children())[:-1]) # ignoring ReLU layer in classifier
# get the features for the whole dataset
features = compute_features(dataloader, model, len(dataset), device) # ndarray, (60k, 512) [-0.019, 0.016]
# cluster the features
clustering_loss = deepcluster.cluster(features, verbose=args.verbose)
# assign pseudo-labels
train_dataset = clustering.cluster_assign(deepcluster.images_lists,
dataset.imgs)
# uniformely sample per target
sampler = UnifLabelSampler(int(args.reassign * len(train_dataset)),
deepcluster.images_lists)
train_dataloader = torch.utils.data.DataLoader(
train_dataset,
batch_size=args.batch,
num_workers=args.workers,
sampler=sampler,
pin_memory=True,
)
# set last fully connected layer
mlp = list(model.classifier.children())
mlp.append(nn.ReLU(inplace=True).to(device))
model.classifier = nn.Sequential(*mlp)
model.top_layer = nn.Linear(fd, len(deepcluster.images_lists))
model.top_layer.weight.data.normal_(0, 0.01)
model.top_layer.bias.data.zero_()
model.top_layer.to(device)
# train network with clusters as pseudo-labels
end = time.time()
# loss = train(train_dataloader, model, criterion, optimizer, epoch)
loss = train(model, device, train_dataloader, optimizer, epoch, criterion)
# print log
if args.verbose:
print('###### Epoch [{0}] ###### \n'
'Time: {1:.3f} s\n'
'Clustering loss: {2:.3f} \n'
'ConvNet loss: {3:.3f}'
.format(epoch, time.time() - end, clustering_loss, loss))
try:
nmi = normalized_mutual_info_score(
clustering.arrange_clustering(deepcluster.images_lists),
clustering.arrange_clustering(cluster_log.data[-1])
)
writer.add_scalar('nmi/train', nmi, epoch)
print('NMI against previous assignment: {0:.3f}'.format(nmi))
except IndexError:
pass
print('####################### \n')
# save running checkpoint
torch.save({'epoch': epoch + 1,
'arch': "VGGMiniCBR",
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict()},
os.path.join(exp_path, 'checkpoint.pth.tar'))
# save cluster assignments
cluster_log.log(deepcluster.images_lists)
torch.save(model.state_dict(), os.path.join(args.exp, "mnist_cnn.pt"))
def main():
global args
args = parser.parse_args()
# fix random seeds
torch.manual_seed(args.seed)
torch.cuda.manual_seed_all(args.seed)
np.random.seed(args.seed)
# CNN
if args.verbose:
print('Architecture: {}'.format(args.arch))
model = models.__dict__[args.arch](sobel=args.sobel)
fd = int(model.top_layer.weight.size()[1])
model.top_layer = None
model.features = torch.nn.DataParallel(model.features)
model.cuda()
cudnn.benchmark = True
# create optimizer
optimizer = torch.optim.SGD(
filter(lambda x: x.requires_grad, model.parameters()),
lr=args.lr,
momentum=args.momentum,
weight_decay=10**args.wd,
)
# define loss function
criterion = nn.CrossEntropyLoss().cuda()
# 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']
# remove top_layer parameters from checkpoint
for key in checkpoint['state_dict']:
if 'top_layer' in key:
del checkpoint['state_dict'][key]
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))
# creating checkpoint repo
exp_check = os.path.join(args.exp, 'checkpoints')
if not os.path.isdir(exp_check):
os.makedirs(exp_check)
# creating cluster assignments log
cluster_log = Logger(os.path.join(args.exp, 'clusters'))
# preprocessing of data
tra = [transforms.ToTensor(), transforms.Normalize((0.1307, ), (0.3081, ))]
# load the data
end = time.time()
# MNIST-full begin:-------------------------------------------
dataset = datasets.MNIST('./data',
train=True,
download=True,
transform=transforms.Compose(tra))
true_label = dataset.train_labels.cpu().numpy()
# MNIST-full end:-------------------------------------------
# # FMNIST begin:-------------------------------------------
# dataset = datasets.FashionMNIST('./data/fmnist', train=True, download=True,
# transform=transforms.Compose(tra))
# true_label = dataset.train_labels.cpu().numpy()
# # FMNIST end:-------------------------------------------
# # MNIST-test begin:-------------------------------------------
# dataset = datasets.MNIST('./data', train=False, download=True,
# transform=transforms.Compose(tra))
# true_label = dataset.test_labels.cpu().numpy()
# # MNIST-test end:-------------------------------------------
# dataset = datasets.ImageFolder(args.data, transform=transforms.Compose(tra))
# if args.verbose: print('Load dataset: {0:.2f} s'.format(time.time() - end))
dataloader = torch.utils.data.DataLoader(dataset,
batch_size=args.batch,
num_workers=args.workers,
pin_memory=True)
# clustering algorithm to use
deepcluster = clustering.__dict__[args.clustering](args.nmb_cluster)
# training convnet with DeepCluster
for epoch in range(args.start_epoch, args.epochs):
end = time.time()
# remove head
model.top_layer = None
model.classifier = nn.Sequential(
*list(model.classifier.children())[:-1])
# get the features for the whole dataset
features = compute_features(dataloader, model, len(dataset))
# cluster the features
clustering_loss = deepcluster.cluster(features, verbose=args.verbose)
# assign pseudo-labels
# train_dataset = clustering.cluster_assign(deepcluster.images_lists,
# dataset.train_data)
train_dataset = clustering.cluster_assign(deepcluster.images_lists,
dataset.train_data)
# uniformely sample per target
sampler = UnifLabelSampler(int(args.reassign * len(train_dataset)),
deepcluster.images_lists)
train_dataloader = torch.utils.data.DataLoader(
train_dataset,
batch_size=args.batch,
num_workers=args.workers,
sampler=sampler,
pin_memory=True,
)
# set last fully connected layer
mlp = list(model.classifier.children())
mlp.append(nn.ReLU(inplace=True).cuda())
model.classifier = nn.Sequential(*mlp)
model.top_layer = nn.Linear(fd, len(deepcluster.images_lists))
model.top_layer.weight.data.normal_(0, 0.01)
model.top_layer.bias.data.zero_()
model.top_layer.cuda()
# train network with clusters as pseudo-labels
end = time.time()
loss = train(train_dataloader, model, criterion, optimizer, epoch)
# print log
if args.verbose:
# print('###### Epoch [{0}] ###### \n'
# 'Time: {1:.3f} s\n'
# 'Clustering loss: {2:.3f} \n'
# 'ConvNet loss: {3:.3f}'
# .format(epoch, time.time() - end, clustering_loss, loss))
try:
y_pred = clustering.arrange_clustering(
deepcluster.images_lists)
y_last = clustering.arrange_clustering(cluster_log.data[-1])
import metrics
acc = metrics.acc(y_pred, y_last)
nmi = metrics.nmi(y_pred, y_last)
acc_ = metrics.acc(true_label, y_pred)
nmi_ = metrics.nmi(true_label, y_pred)
print(
'ACC=%.4f, NMI=%.4f; Relative ACC=%.4f, Relative NMI=%.4f'
% (acc_, nmi_, acc, nmi))
except IndexError:
pass
print('####################### \n')
# save running checkpoint
torch.save(
{
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict()
}, os.path.join(args.exp, 'checkpoint.pth.tar'))
# save cluster assignments
cluster_log.log(deepcluster.images_lists)
def main(args):
# fix random seeds
torch.manual_seed(args.seed)
torch.cuda.manual_seed_all(args.seed)
np.random.seed(args.seed)
device = torch.device('cuda:0' if torch.cuda.is_available() else "cpu")
print(device)
# CNN
if args.verbose:
print('Architecture: {}'.format(args.arch))
model = models.__dict__[args.arch](sobel=False,
bn=True,
out=args.nmb_cluster)
fd = int(model.top_layer[0].weight.size()
[1]) # due to transpose, fd is input dim of W (in dim, out dim)
model.top_layer = None
model.features = torch.nn.DataParallel(model.features)
model = model.double()
model.to(device)
cudnn.benchmark = True
if args.optimizer is 'Adam':
print('Adam optimizer: conv')
optimizer = torch.optim.Adam(
filter(lambda x: x.requires_grad, model.parameters()),
lr=args.lr_Adam,
betas=(0.5, 0.99),
weight_decay=10**args.wd,
)
else:
print('SGD optimizer: conv')
optimizer = torch.optim.SGD(
filter(lambda x: x.requires_grad, model.parameters()),
lr=args.lr_SGD,
momentum=args.momentum,
weight_decay=10**args.wd,
)
criterion = nn.CrossEntropyLoss()
# 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']
# remove top_layer parameters from checkpoint
copy_checkpoint_state_dict = checkpoint['state_dict'].copy()
for key in list(copy_checkpoint_state_dict):
if 'top_layer' in key:
del copy_checkpoint_state_dict[key]
checkpoint['state_dict'] = copy_checkpoint_state_dict
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))
# creating checkpoint repo
exp_check = os.path.join(args.exp, '../../..', 'checkpoints')
if not os.path.isdir(exp_check):
os.makedirs(exp_check)
# creating cluster assignments log
cluster_log = Logger(os.path.join(args.exp, '../../..', 'clusters.pickle'))
# # Create echogram sampling index
print('Sample echograms.')
end = time.time()
dataset_cp = sampling_echograms_full(args)
dataloader_cp = torch.utils.data.DataLoader(dataset_cp,
shuffle=False,
batch_size=args.batch,
num_workers=args.workers,
drop_last=False,
pin_memory=True)
if args.verbose:
print('Load dataset: {0:.2f} s'.format(time.time() - end))
# clustering algorithm to use
deepcluster = clustering.__dict__[args.clustering](args.nmb_cluster,
args.pca)
# deepcluster = clustering.Kmeans(no.cluster, dim.pca)
loss_collect = [[], [], [], [], []]
nmi_save = []
# for evaluation
dataset_eval = sampling_echograms_eval(args)
eval_dataloader = torch.utils.data.DataLoader(
dataset_eval,
batch_size=args.batch,
shuffle=False,
num_workers=args.workers,
pin_memory=True,
)
# training convnet with DeepCluster
for epoch in range(args.start_epoch, args.epochs):
# remove head
model.top_layer = None
model.classifier = nn.Sequential(*list(model.classifier.children()))
# get the features for the whole dataset
features_train, input_tensors_train, labels_train = compute_features(
dataloader_cp, model, len(dataset_cp), device=device, args=args)
# cluster the features
print('Cluster the features')
end = time.time()
clustering_loss, pca_features = deepcluster.cluster(
features_train, verbose=args.verbose)
# deepcluster.cluster(features_train, verbose=args.verbose)
print('Cluster time: {0:.2f} s'.format(time.time() - end))
nan_location = np.isnan(pca_features)
inf_location = np.isinf(pca_features)
if (not np.allclose(nan_location, 0)) or (not np.allclose(
inf_location, 0)):
print('PCA: Feature NaN or Inf found. Nan count: ',
np.sum(nan_location), ' Inf count: ', np.sum(inf_location))
print('Skip epoch ', epoch)
torch.save(pca_features, 'pca_NaN_%d.pth.tar' % epoch)
torch.save(features_train, 'feature_NaN_%d.pth.tar' % epoch)
continue
# save patches per epochs
cp_epoch_out = [
features_train, deepcluster.images_lists,
deepcluster.images_dist_lists, input_tensors_train, labels_train
]
linear_svc = SimpleClassifier(epoch,
cp_epoch_out,
tr_size=5,
iteration=20)
if args.verbose:
print('###### Epoch [{0}] ###### \n'
'Classify. accu.: {1:.3f} \n'
'Pairwise classify. accu: {2} \n'.format(
epoch, linear_svc.whole_score, linear_svc.pair_score))
if (epoch % args.save_epoch == 0):
end = time.time()
with open(
os.path.join(args.exp, '../../..',
'cp_epoch_%d.pickle' % epoch), "wb") as f:
pickle.dump(cp_epoch_out, f)
with open(
os.path.join(args.exp, '../../..',
'pca_epoch_%d.pickle' % epoch), "wb") as f:
pickle.dump(pca_features, f)
print('Feature save time: {0:.2f} s'.format(time.time() - end))
# assign pseudo-labels
print('Assign pseudo labels')
size_cluster = np.zeros(len(deepcluster.images_lists))
for i, _list in enumerate(deepcluster.images_lists):
size_cluster[i] = len(_list)
print('size in clusters: ', size_cluster)
img_label_pair_train = zip_img_label(input_tensors_train, labels_train)
train_dataset = clustering.cluster_assign(
deepcluster.images_lists,
img_label_pair_train) # Reassigned pseudolabel
# uniformly sample per target
sampler_train = UnifLabelSampler(int(len(train_dataset)),
deepcluster.images_lists)
train_dataloader = torch.utils.data.DataLoader(
train_dataset,
batch_size=args.batch,
shuffle=False,
num_workers=args.workers,
sampler=sampler_train,
pin_memory=True,
)
# set last fully connected layer
mlp = list(model.classifier.children()
) # classifier that ends with linear(512 * 128)
mlp.append(nn.ReLU().to(device))
model.classifier = nn.Sequential(*mlp)
model.top_layer = nn.Sequential(
nn.Linear(fd, args.nmb_cluster),
nn.Softmax(dim=1),
)
# model.top_layer = nn.Linear(fd, args.nmb_cluster)
model.top_layer[0].weight.data.normal_(0, 0.01)
model.top_layer[0].bias.data.zero_()
model.top_layer = model.top_layer.double()
model.top_layer.to(device)
# train network with clusters as pseudo-labels
end = time.time()
with torch.autograd.set_detect_anomaly(True):
# loss, tr_epoch_out = train(train_dataloader, model, criterion, optimizer, epoch, device=device, args=args)
loss = train(train_dataloader,
model,
criterion,
optimizer,
epoch,
device=device,
args=args)
print('Train time: {0:.2f} s'.format(time.time() - end))
# if (epoch % args.save_epoch == 0):
# end = time.time()
# with open(os.path.join(args.exp, '..', 'tr_epoch_%d.pickle' % epoch), "wb") as f:
# pickle.dump(tr_epoch_out, f)
# print('Save train time: {0:.2f} s'.format(time.time() - end))
# Accuracy with training set (output vs. pseudo label)
# accuracy_tr = np.mean(tr_epoch_out[1] == np.argmax(tr_epoch_out[2], axis=1))
# print log
if args.verbose:
print('###### Epoch [{0}] ###### \n'
'Time: {1:.3f} s\n'
'ConvNet tr_loss: {2:.3f} \n'
'Clustering loss: {3:.3f} \n'.format(epoch,
time.time() - end, loss,
clustering_loss))
try:
nmi = normalized_mutual_info_score(
clustering.arrange_clustering(deepcluster.images_lists),
clustering.arrange_clustering(cluster_log.data[-1]))
nmi_save.append(nmi)
print('NMI against previous assignment: {0:.3f}'.format(nmi))
with open("./nmi_collect.pickle", "wb") as ff:
pickle.dump(nmi_save, ff)
except IndexError:
pass
print('####################### \n')
# save running checkpoint
torch.save(
{
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict()
}, os.path.join(args.exp, '../../..', 'checkpoint.pth.tar'))
# evaluation: echogram reconstruction
if (epoch % args.save_epoch == 0):
eval_epoch_out = evaluate(eval_dataloader,
model,
device=device,
args=args)
with open(
os.path.join(args.exp, '../../..',
'eval_epoch_%d.pickle' % epoch), "wb") as f:
pickle.dump(eval_epoch_out, f)
print('epoch: ', type(epoch), epoch)
print('loss: ', type(loss), loss)
print('linear_svc.whole_score: ', type(linear_svc.whole_score),
linear_svc.whole_score)
print('linear_svc.pair_score: ', type(linear_svc.pair_score),
linear_svc.pair_score)
print('clustering_loss: ', type(clustering_loss), clustering_loss)
loss_collect[0].append(epoch)
loss_collect[1].append(loss)
loss_collect[2].append(linear_svc.whole_score)
loss_collect[3].append(linear_svc.pair_score)
loss_collect[4].append(clustering_loss)
with open(os.path.join(args.exp, '../../..', 'loss_collect.pickle'),
"wb") as f:
pickle.dump(loss_collect, f)
# save cluster assignments
cluster_log.log(deepcluster.images_lists)
def main(args):
# fix random seeds
seed(31)
# CNN
model = MobileNetV1(num_classes=100, sobel=True)
fd = int(model.top_layer.weight.size()[1])
model.top_layer = None
model.features = torch.nn.DataParallel(model.features)
model.cuda()
cudnn.benchmark = True
# create optimizer
optimizer = torch.optim.SGD(
[x for x in model.parameters() if x.requires_grad],
lr=args.lr,
momentum=args.momentum,
weight_decay=10**args.wd,
)
# define loss function
criterion = nn.CrossEntropyLoss().cuda()
# creating cluster assignments log
cluster_log = Logger(os.path.join('./image_list_log/', 'clusters'))
end = time.time()
# load the data
dataset = datasets.ImageFolder(root=r'./dataset/train',
transform=transform())
dataloader = torch.utils.data.DataLoader(dataset,
batch_size=args.batch,
num_workers=args.workers,
pin_memory=True)
# clustering algorithm to use
deepcluster = clustering.__dict__[args.clustering](args.nmb_cluster)
print('start train')
# training convnet with DeepCluster
for epoch in range(args.start_epoch, args.epochs):
print(epoch)
# remove head
model.top_layer = None
model.classifier = nn.Sequential(
*list(model.classifier.children())[:-1])
# get the features for the whole dataset
features = compute_features(dataloader, model, len(dataset),
args.batch)
# cluster the feature
clustering_loss = deepcluster.cluster(features)
# assign pseudo-labels
train_dataset = clustering.cluster_assign(deepcluster.images_lists,
dataset.imgs)
# uniformly sample per target
sampler = UnifLabelSampler(int(args.reassign * len(train_dataset)),
deepcluster.images_lists)
train_dataloader = torch.utils.data.DataLoader(
train_dataset,
batch_size=args.batch,
num_workers=args.workers,
sampler=sampler,
pin_memory=True,
)
# set last fully connected layer
mlp = list(model.classifier.children())
mlp.append(nn.ReLU(inplace=True).cuda())
model.classifier = nn.Sequential(*mlp)
model.top_layer = nn.Linear(fd, len(deepcluster.images_lists))
model.top_layer.weight.data.normal_(0, 0.01)
model.top_layer.bias.data.zero_()
model.top_layer.cuda()
# train network with clusters as pseudo-labels
end = time.time()
loss = train(train_dataloader, model, criterion, optimizer, epoch,
args.lr, args.wd)
# print log
# print('###### Epoch [{0}] ###### \n'
# 'Time: {1:.3f} s\n'
# 'Clustering loss: {2:.3f} \n'
# 'ConvNet loss: {3:.3f}'
# .format(epoch, time.time() - end, clustering_loss, loss))
try:
nmi = normalized_mutual_info_score(
clustering.arrange_clustering(deepcluster.images_lists),
clustering.arrange_clustering(cluster_log.data[-1]))
print('NMI against previous assignment: {0:.3f}'.format(nmi))
f = open('result.txt', "a")
f.write('NMI against previous assignment: {0:.3f}'.format(nmi))
f.close()
# print(loss)
except IndexError:
pass
print('####################### \n')
# save cluster assignments
cluster_log.log(deepcluster.images_lists)
