def print_tensor(x: torch.tensor, pt=True, val=True, shp=True):
"""Print the mean, min, median, std, and size of a tensor tensor
Args:
x:
val: if print the values of the tensor
shp: if print the shape of the tensor
Returns: None
"""
x = x.float()
message = ''
if shp:
message = str(x.shape) + '\n'
if val:
x = x.flatten()
if len(x) != 1:
message += (
'mean = %3.3f, min = %3.3f, max = %3.3f, median = %3.3f, std=%3.3f'
% (x.mean(), x.min(), x.max(), x.median(), x.std()))
else:
message += (f'one element {x[0]}')
if pt:
logging.debug(message)
return message
def forward(self, x: torch.tensor) -> torch.tensor:
"""
Forward pass function
Args:
x (torch.tensor): Input tensors of dimensions [B, 2, 14, 14] with B being batch size
Returns:
torch.tensor: Predicted probability of size [1]
"""
x = self.bn2d(self.conv1(x))
x = self.relu(x)
x = self.drop2d(self.conv2(x))
x = self.relu(self.pool(x))
x = self.drop(self.bn(self.fc1(x.flatten(start_dim=1))))
x = self.relu(x)
x = self.drop(self.fc2(x))
x = self.relu(x)
x = self.drop(self.fc3(x))
x = self.relu(x)
x = self.sigmoid(self.fc4(x))
return x.squeeze(), None
def forward(ctx, inputs: torch.tensor, threshold: float, sigmoid: bool):
"""
Args:
inputs (`torch.FloatTensor`)
The input matrix from which the binarizer computes the binary mask.
threshold (`float`)
The threshold value (in R).
sigmoid (`bool`)
If set to ``True``, we apply the sigmoid function to the `inputs` matrix before comparing to `threshold`.
In this case, `threshold` should be a value between 0 and 1.
Returns:
mask (`torch.FloatTensor`)
Binary matrix of the same size as `inputs` acting as a mask (1 - the associated weight is
retained, 0 - the associated weight is pruned).
"""
nb_elems = inputs.numel()
nb_min = int(0.005 * nb_elems) + 1
if sigmoid:
mask = (torch.sigmoid(inputs) > threshold).type(inputs.type())
else:
mask = (inputs > threshold).type(inputs.type())
if mask.sum() < nb_min:
# We limit the pruning so that at least 0.5% (half a percent) of the weights are remaining
k_threshold = inputs.flatten().kthvalue(max(nb_elems - nb_min,
1)).values
mask = (inputs > k_threshold).type(inputs.type())
return mask
def forward_once(self, x: torch.tensor) -> torch.tensor:
"""
Forward pass function for the global siamese CNN
Args:
x [float32]: input images with dimension Bx2x14x14 (for batch size B)
Returns:
[int]: predicted probability ]0,1[
[float32] : predicted classe by pair, size Bx2x10
"""
# flatten image input
x = x.flatten(start_dim=1) # (-1, 2x14x14)
# add hidden layer, with relu activation function
x = self.relu(self.fc1(x))
x = self.drop(x)
x = self.relu(self.fc2(x))
x = self.drop(x)
x = self.relu(self.fc3(x))
x = self.drop(x)
x = self.fc4(x)
return x
def forward(ctx, inputs: torch.tensor, threshold: float, sigmoid: bool):
"""
Args:
inputs (`torch.FloatTensor`)
The input matrix from which the binarizer computes the binary mask.
threshold (`float`)
The percentage of weights to keep (the rest is pruned).
`threshold` is a float between 0 and 1.
sigmoid (`bool`)
Whether to apply a sigmoid on the threshold
Returns:
mask (`torch.FloatTensor`)
Binary matrix of the same size as `inputs` acting as a mask (1 - the associated weight is
retained, 0 - the associated weight is pruned).
"""
# Get the subnetwork by sorting the inputs and using the top threshold
if sigmoid:
threshold = torch.sigmoid(threshold).item()
ctx.sigmoid = sigmoid
mask = inputs.clone()
_, idx = inputs.flatten().sort(descending=True)
j = math.ceil(threshold * inputs.numel())
# flat_out and mask access the same memory.
flat_out = mask.flatten()
flat_out[idx[j:]] = 0.
flat_out[idx[:j]] = 1.
ctx.save_for_backward(mask)
return mask
def forward(ctx, inputs: torch.tensor, threshold: float):
"""
Args:
inputs (`torch.FloatTensor`)
The input matrix from which the binarizer computes the binary mask.
threshold (`float`)
The percentage of weights to keep (the rest is pruned).
`threshold` is a float between 0 and 1.
Returns:
mask (`torch.FloatTensor`)
Binary matrix of the same size as `inputs` acting as a mask (1 - the associated weight is
retained, 0 - the associated weight is pruned).
"""
# Get the subnetwork by sorting the inputs and using the top threshold %
if not isinstance(threshold, float):
threshold = threshold[0]
mask = inputs.clone()
_, idx = inputs.flatten().sort(descending=True)
j = int(threshold * inputs.numel())
# flat_out and mask access the same memory.
flat_out = mask.flatten()
flat_out[idx[j:]] = 0
flat_out[idx[:j]] = 1
return mask
def predict(self, Xnew: torch.tensor, Xtrain: torch.tensor,
ytrain: torch.tensor):
"""
Inputs:
:Xnew: n_newsamples * n_features
:Xtrain: n_samples * n_features
:ytrain: n_samples / n_samples * n_output
Returns:
A set of 1d normal distributions, their mean/var are scalar tensors
"""
device = Xnew.device
Xnew, Xtrain = self.transformer(Xnew), self.transformer(Xtrain)
K = self.to_matrix(self.kernel(Xtrain))
Kprime = K + self.alpha * torch.eye(K.shape[0], device=device)
L = torch.cholesky(Kprime)
kXnew = self.to_matrix(self.kernel(Xnew, Xtrain))
ytrain = ytrain.flatten()
ytrain_minus_mean = ytrain - self.mean(Xtrain)
# logging.debug(self.mean)
mean_ynew = self.mean(Xnew) + torch.squeeze(torch.matmul(
kXnew, torch.cholesky_solve(ytrain_minus_mean[:, None], L)),
dim=-1)
logging.debug("mean={}".format(mean_ynew))
# LL^T _x = kXnew^T
_x = torch.cholesky_solve(kXnew.T, L)
kXnewXnew = self.to_matrix(self.kernel(Xnew))
std_ynew = torch.sqrt(
torch.diag(kXnewXnew) - torch.einsum("ij,ji->i", kXnew, _x))
logging.debug("var={}".format(std_ynew))
return [Normal(mu, sigma) for mu, sigma in zip(mean_ynew, std_ynew)]
def forward_once(self, x: torch.tensor) -> torch.tensor:
"""
Forward pass function used in the sub-network
Args:
x [float32]: input image with dimension Bx1x14x14 (for batch size B)
Returns:
[float32]: non activated tensor of dimension Bx1x10
"""
# flatten image input
x = x.flatten(start_dim=1)
# add hidden layer, with relu activation function
x = self.relu(self.fc1(x))
x = self.drop(x)
x = self.relu(self.fc2(x))
x = self.drop(x)
x = self.relu(self.fc3(x))
x = self.drop(x)
x = self.fc4(x)
return x
def com_reduce(buffer_m: torch.tensor, rank, world_size, comm):
tensor_size = torch.numel(buffer_m)
chunk_size = (tensor_size + world_size - 1) // world_size
last_chunk_size = tensor_size - chunk_size * (world_size - 1)
my_chunk_size = last_chunk_size if rank == world_size - 1 else chunk_size
flatten_buffer_m = buffer_m.flatten()
flatten_buffer_m_cupy = cupy.fromDlpack(to_dlpack(flatten_buffer_m))
# First round of communication
recvbuf = cupy.zeros([world_size, my_chunk_size],
dtype=flatten_buffer_m_cupy.dtype)
requests = []
for idx in range(world_size):
start = idx * chunk_size
length = last_chunk_size if idx == world_size - 1 else chunk_size
req_sign = comm.Igather(flatten_buffer_m_cupy[start:start + length],
recvbuf,
root=idx)
requests.append(req_sign)
MPI.Request.Waitall(requests)
# Second round of communication
recvbuf_flatten = recvbuf.flatten()
local_reduced_chunk = cupy.zeros(my_chunk_size,
dtype=flatten_buffer_m_cupy.dtype)
_avg_chunks(recvbuf_flatten, my_chunk_size, world_size,
local_reduced_chunk)
recvbuf_server = [
cupy.zeros(chunk_size, dtype=flatten_buffer_m_cupy.dtype)
] * (world_size - 1)
recvbuf_server.append(
cupy.zeros(last_chunk_size, dtype=flatten_buffer_m_cupy.dtype))
recvbuf_server[rank] = local_reduced_chunk
server_requests = []
for idx in range(world_size):
if idx != rank:
req_server_send = comm.Isend(local_reduced_chunk, idx)
req_server_recv = comm.Irecv(recvbuf_server[idx], idx)
server_requests.append(req_server_send)
server_requests.append(req_server_recv)
MPI.Request.Waitall(server_requests)
recvbuf_server_flatten = cupy.concatenate(recvbuf_server)
aggregated_m_tensor = from_dlpack(recvbuf_server_flatten.toDlpack())
buffer_m.set_(aggregated_m_tensor.type(buffer_m.dtype).view_as(buffer_m))
def to_sparse_by_cdf(t: torch.tensor, cdf: float):
N = t.size(0)
t_flatten = t.flatten(start_dim=1).clone().detach()
sorted_t_flatten, indices = torch.sort(t_flatten, descending=True, dim=-1)
sorted_t_flatten_cum = torch.cumsum(sorted_t_flatten, dim=-1)
for i in range(N):
mask = sorted_t_flatten_cum[i] < cdf
mask[torch.sum(mask)] = True
t_flatten[i, indices[i, mask]] = 1
t_flatten[i, indices[i, ~mask]] = 0
return t_flatten.reshape(*t.size()).long()
def add(self, state: tensor, action: tensor, reward: tensor, done: bool,
new_state: tensor):
if self.cur_buffer_size >= self.max_buffer_size:
del self.states[0]
del self.actions[0]
del self.rewards[0]
del self.dones[0]
del self.new_states[0]
self.states.append(state)
self.actions.append(action.flatten())
self.rewards.append(reward)
self.dones.append(done)
self.new_states.append(new_state)
self.cur_buffer_size = len(self.states)
def forward_once(self, x: torch.tensor) -> torch.tensor:
"""
Forward pass function used in the sub-network
Args:
x [float32]: input image with dimension Bx1x14x14 (for batch size B)
Returns:
[float32]: non activated tensor of dimension Bx1x10
"""
x = self.drop(self.bn2d(self.conv1(x)))
x = self.relu(x)
x = self.drop2d(self.conv2(x))
x = self.relu(self.pool(x))
x = self.drop(self.bn(self.fc1(x.flatten(start_dim=1))))
x = self.relu(x)
x = self.drop(self.fc2(x))
return x