def gen_testloss(args):
# load data and model
params = utils.load_params(args.model_dir)
ckpt_dir = os.path.join(args.model_dir, 'checkpoints')
ckpt_paths = [int(e[11:-3]) for e in os.listdir(ckpt_dir) if e[-3:] == ".pt"]
ckpt_paths = np.sort(ckpt_paths)
# csv
headers = ["epoch", "step", "loss", "discrimn_loss_e", "compress_loss_e",
"discrimn_loss_t", "compress_loss_t"]
csv_path = utils.create_csv(args.model_dir, 'losses_test.csv', headers)
print('writing to:', csv_path)
# load data
test_transforms = tf.load_transforms('test')
testset = tf.load_trainset(params['data'], test_transforms, train=False)
testloader = DataLoader(testset, batch_size=params['bs'], shuffle=False, num_workers=4)
# save loss
criterion = MaximalCodingRateReduction(gam1=params['gam1'], gam2=params['gam2'], eps=params['eps'])
for epoch, ckpt_path in enumerate(ckpt_paths):
net, epoch = tf.load_checkpoint(args.model_dir, epoch=epoch, eval_=True)
for step, (batch_imgs, batch_lbls) in enumerate(testloader):
features = net(batch_imgs.cuda())
loss, loss_empi, loss_theo = criterion(features, batch_lbls,
num_classes=len(testset.num_classes))
utils.save_state(args.model_dir, epoch, step, loss.item(),
*loss_empi, *loss_theo, filename='losses_test.csv')
print("Finished generating test loss.")
## Prepare for Training
if args.pretrain_dir is not None:
net, _ = tf.load_checkpoint(args.pretrain_dir, args.pretrain_epo)
utils.update_params(model_dir, args.pretrain_dir)
else:
net = tf.load_architectures(args.arch, args.fd)
transforms = tf.load_transforms(args.transform)
trainset = tf.load_trainset(args.data, path=args.data_dir)
trainloader = AugmentLoader(trainset,
transforms=transforms,
sampler=args.sampler,
batch_size=args.bs,
num_aug=args.aug)
criterion = MaximalCodingRateReduction(gam1=args.gam1,
gam2=args.gam2,
eps=args.eps)
optimizer = optim.SGD(net.parameters(),
lr=args.lr,
momentum=args.mom,
weight_decay=args.wd)
scheduler = lr_scheduler.MultiStepLR(optimizer, [30, 60], gamma=0.1)
utils.save_params(model_dir, vars(args))
## Training
for epoch in range(args.epo):
for step, (batch_imgs, _, batch_idx) in enumerate(trainloader):
batch_features = net(batch_imgs.cuda())
loss, loss_empi, loss_theo = criterion(batch_features, batch_idx)
optimizer.zero_grad()
loss.backward()
import torch.nn.functional as F
from torch.optim import SGD, lr_scheduler
from sklearn.metrics.cluster import normalized_mutual_info_score as nmi_score
from sklearn.metrics import adjusted_rand_score as ari_score
from sklearn.cluster import KMeans
from utils.util import BCE, PairEnum, cluster_acc, Identity, AverageMeter, seed_torch
from utils import ramps
from models.resnet import ResNet, BasicBlock
from data.cifarloader import CIFAR10Loader, CIFAR10LoaderMix, CIFAR100Loader, CIFAR100LoaderMix
from data.svhnloader import SVHNLoader, SVHNLoaderMix
from tqdm import tqdm
import numpy as np
import os
from loss import MaximalCodingRateReduction
MCR = MaximalCodingRateReduction(gam1=20., gam2=0.5, eps=0.5)
def smooth_loss(feat,mask_lb):
feature = (feat[~mask_lb])
feature = F.normalize(feature,1)
inputs = feature.mm(feature.t())
targets = smooth_hot(inputs.detach().clone())
outputs = F.log_softmax(inputs, dim=1)
loss = - (targets * outputs)
loss = loss.sum(dim=1)
loss = loss.mean(dim=0)
return loss
def smooth_hot(inputs, k=6):