def opt_sk(model, selflabels_in, epoch):
if args.hc == 1:
PS = np.zeros((len(trainloader.dataset), args.ncl))
else:
PS_pre = np.zeros((len(trainloader.dataset), knn_dim))
for batch_idx, (data, _, _selected) in enumerate(trainloader):
data = data.cuda()
if args.hc == 1:
p = nn.functional.softmax(model(data), 1)
PS[_selected, :] = p.detach().cpu().numpy()
else:
p = model(data)
PS_pre[_selected, :] = p.detach().cpu().numpy()
if args.hc == 1:
selflabels = optimize_L_sk(PS)
else:
_nmis = np.zeros(args.hc)
nh = epoch % args.hc # np.random.randint(args.hc)
logger.info("computing head {}".format(nh))
tl = getattr(model, "top_layer{}".format(nh))
# do the forward pass:
PS = (PS_pre @ tl.weight.cpu().numpy().T + tl.bias.cpu().numpy())
PS = py_softmax(PS, 1)
selflabels_ = optimize_L_sk(PS)
selflabels_in[nh] = selflabels_
selflabels = selflabels_in
return selflabels
def opt_sk(model, selflabels_in, epoch):
if args.hc == 1:
PS = np.zeros((len(trainloader.dataset), args.ncl))
else:
PS_pre = np.zeros((len(trainloader.dataset), knn_dim))
for batch_idx, (data, _, _selected) in enumerate(trainloader):
data = data.cuda()
if args.hc == 1:
p = nn.functional.softmax(model(data), 1)
PS[_selected, :] = p.detach().cpu().numpy()
else:
p = model(data)
PS_pre[_selected, :] = p.detach().cpu().numpy()
if args.hc == 1:
cost, selflabels = optimize_L_sk(PS)
_costs = [cost]
else:
_nmis = np.zeros(args.hc)
_costs = np.zeros(args.hc)
nh = epoch % args.hc # np.random.randint(args.hc)
print("computing head %s " % nh, end="\r", flush=True)
tl = getattr(model, "top_layer%d" % nh)
# do the forward pass:
PS = (PS_pre @ tl.weight.cpu().numpy().T + tl.bias.cpu().numpy())
PS = py_softmax(PS, 1)
c, selflabels_ = optimize_L_sk(PS)
_costs[nh] = c
selflabels_in[nh] = selflabels_
selflabels = selflabels_in
return selflabels
def cpu_sk(self):
""" Sinkhorn Knopp optimization on CPU
* stores activations to RAM
* does matrix-vector multiplies on CPU
* slower than GPU
"""
# 1. aggregate inputs:
N = len(self.pseudo_loader.dataset)
if self.num_heads == 1:
self.PS = np.zeros((N, self.num_clusters_per_head), dtype=self.dtype)
else:
self.PS_pre = np.zeros((N, self.presize), dtype=self.dtype)
now = time.time()
l_dl = len(self.pseudo_loader)
time.time()
batch_time = MovingAverage(intertia=0.9)
self.model.headcount = 1
for batch_idx, (data, _, _selected) in enumerate(self.pseudo_loader):
data = data.to(self.device)
mass = data.size(0)
if self.num_heads == 1:
p = nn.functional.softmax(self.model(data), 1)
self.PS[_selected, :] = p.detach().cpu().numpy().astype(self.dtype)
else:
p = self.model(data)
self.PS_pre[_selected, :] = p.detach().cpu().numpy().astype(self.dtype)
batch_time.update(time.time() - now)
now = time.time()
if batch_idx % 50 == 0:
print(f"Aggregating batch {batch_idx:03}/{l_dl}, speed: {mass / batch_time.avg:04.1f}Hz",
end='\r', flush=True)
self.model.headcount = self.num_heads
print("Aggreg of outputs took {0:.2f} min".format((time.time() - now) / 60.), flush=True)
# 2. solve label assignment via sinkhorn-knopp:
if self.num_heads == 1:
optimize_L_sk(self, nh=0)
else:
for nh in range(self.num_heads):
print(f"computing head {nh} ", end="\r", flush=True)
tl = getattr(self.model, f"top_layer{nh:d}")
time_mat = time.time()
# clear memory
try:
del self.PS
except:
pass
# apply last FC layer (a matmul and adding of bias)
self.PS = (self.PS_pre @ tl.weight.cpu().numpy().T.astype(self.dtype)
+ tl.bias.cpu().numpy().astype(self.dtype))
print(f"matmul took {(time.time() - time_mat) / 60:.2f}min", flush=True)
self.PS = py_softmax(self.PS, 1)
optimize_L_sk(self, nh=nh)
return
