Replacing a for loop with indexing:kmeans

username: abd
cooked: <p>Hi,<br>
I’m trying to replace the "for loop " with indexing.</p>
<p>This the code :<br>
dataset=torch.randn(200,1000,64).cuda()<br>
t=torch.zeros(200,4,64)<br>
for i in range(200):<br>
centers, codes = cluster(dataset[i],4)<br>
t[i]=centers</p>

username: ptrblck
cooked: <p>Could you post the definition of <code>cluster</code>?<br>
If you want to replace the loop, <code>cluster</code> would have to process batched data.</p>

username: abd
cooked: <p>Hi ptrblck,<br>
How can I make the function cluster process batched data?The code of cluster is as follows:<br>
import torch<br>
import numpy as np<br>
import random<br>
import sys<br>
import timeit<br>
device_gpu = torch.device(‘cuda’)</p>
<h1>Choosing <code>num_centers</code> random data points as the initial centers</h1>
<p>def random_init(dataset, num_centers):<br>
num_points = dataset.size(0)<br>
dimension = dataset.size(1)<br>
used = torch.zeros(num_points, dtype=torch.long)<br>
indices = torch.zeros(num_centers, dtype=torch.long)<br>
for i in range(num_centers):<br>
while True:<br>
cur_id = random.randint(0, num_points - 1)<br>
if used[cur_id] &gt; 0:<br>
continue<br>
used[cur_id] = 1<br>
indices[i] = cur_id<br>
break<br>
indices = indices.to(device_gpu)<br>
centers = torch.gather(dataset, 0, indices.view(-1, 1).expand(-1, dimension))<br>
return centers</p>
<h1>Compute for each data point the closest center</h1>
<p>def compute_codes(dataset, centers):<br>
num_points = dataset.size(0)<br>
dimension = dataset.size(1)<br>
num_centers = centers.size(0)<br>
# 5e8 should vary depending on the free memory on the GPU<br>
# Ideally, automatically <img src="https://discuss.pytorch.org/images/emoji/apple/wink.png?v=9" title=":wink:" class="emoji" alt=":wink:"><br>
chunk_size = int(5e8 / num_centers)<br>
codes = torch.zeros(num_points, dtype=torch.long, device=device_gpu)<br>
centers_t = torch.transpose(centers, 0, 1)<br>
centers_norms = torch.sum(centers ** 2, dim=1).view(1, -1)<br>
for i in range(0, num_points, chunk_size):<br>
begin = i<br>
end = min(begin + chunk_size, num_points)<br>
dataset_piece = dataset[begin:end, :]<br>
dataset_norms = torch.sum(dataset_piece ** 2, dim=1).view(-1, 1)<br>
distances = torch.mm(dataset_piece, centers_t)<br>
distances *= -2.0<br>
distances += dataset_norms<br>
distances += centers_norms<br>
_, min_ind = torch.min(distances, dim=1)<br>
codes[begin:end] = min_ind<br>
return codes</p>
<h1>Compute new centers as means of the data points forming the clusters</h1>
<p>def update_centers(dataset, codes, num_centers):<br>
num_points = dataset.size(0)<br>
dimension = dataset.size(1)<br>
centers = torch.zeros(num_centers, dimension, dtype=torch.float, device=device_gpu)<br>
cnt = torch.zeros(num_centers, dtype=torch.float, device=device_gpu)<br>
centers.scatter_add_(0, codes.view(-1, 1).expand(-1, dimension), dataset)<br>
cnt.scatter_add_(0, codes, torch.ones(num_points, dtype=torch.float, device=device_gpu))<br>
# Avoiding division by zero<br>
# Not necessary if there are no duplicates among the data points<br>
cnt = torch.where(cnt &gt; 0.5, cnt, torch.ones(num_centers, dtype=torch.float, device=device_gpu))<br>
centers /= cnt.view(-1, 1)<br>
return centers<br>
<span class="hashtag">#Cluster</span><br>
def cluster(dataset, num_centers):<br>
centers = random_init(dataset, num_centers)<br>
codes = compute_codes(dataset, centers)<br>
num_iterations = 0<br>
while True:<br>
num_iterations += 1<br>
centers = update_centers(dataset, codes, num_centers)<br>
new_codes = compute_codes(dataset, centers)<br>
# Waiting until the clustering stops updating altogether<br>
# This is too strict in practice<br>
if torch.equal(codes, new_codes):<br>
break<br>
codes = new_codes<br>
return centers, codes</p>
<p>if <strong>name</strong> == ‘<strong>main</strong>’:</p>
<pre><code>dataset=torch.randn(200,1000,64).cuda()
t=torch.zeros(200,2,64)
print(dataset.shape)
print('Starting clustering')
start=timeit.default_timer()
for i in range(200):
     centers, codes = cluster(dataset[i], 2)
     t[i]=centers</code></pre>