I defined following function for some simple matrix operations. I found that when cs is fixed, and X_cuda (a (N,D) matrix) is large, function "cnfg = cs_cuda[candidate_ind, :] "consumes the majority of the time, it consumes more time than all rest of the operations. Then I profiled this function, and I found that: aten::_local_scalar_dense: This operation is taking up 97.99% of the CPU time. Is anyone know why this happens def seleval_GPU(cs): # Convert data to PyTorch tensors N,D = X_cuda.shape cs_cuda = cs.cuda() M = cs_cuda.size(0) # Convert indices_list to a tensor A_chunk = X_cuda[cs_cuda].float() #(M, D, D) b_chunk = torch.ones(M, D, 1, device='cuda') ws = torch.linalg.solve(A_chunk, b_chunk) # (M, D, 1) # convert to float16 ws = ws.squeeze(-1) #(M,D) ones = torch.ones(M, 1).cuda() #(M,1) sign_dists= torch.matmul(ws, X_cuda.T)-ones # (M,N) threshold = 1e-6 sign_dists[sign_dists.abs() < threshold] = 0 sign_vectors = torch.sign(sign_dists) # Use advanced indexing to set the elements at the specified indices to 0 # sign_vectors[torch.arange(M).unsqueeze(1), cs_cuda] = 0 sum_labelss = sign_vectors + t_cuda lss_pos = torch.sum(sum_labelss == 0, dim=1, dtype=torch.int16) #size (M,) lss_neg = N-D-lss_pos best_lss = torch.stack((lss_pos, lss_neg), dim=1) best_ls, min_index = torch.min(best_lss,dim=1) candidate_loss, candidate_ind = torch.min(best_ls, dim=0) candidate_ind = candidate_ind.cuda() # Ensure candidate_ind is on GPU start = time.perf_counter() cnfg = cs_cuda[candidate_ind, :] # Accessing large tensor on GPU end = time.perf_counter() print(f'you spend {end-start} time for this block') best_cnfg = (cnfg, candidate_loss) return best_cnfg Continue reading...
