diff --git a/core/NeuralEssentials/cudamodel.py b/core/NeuralEssentials/cudamodel.py index 57c4900..9b14607 100644 --- a/core/NeuralEssentials/cudamodel.py +++ b/core/NeuralEssentials/cudamodel.py @@ -34,10 +34,10 @@ def check_precision_device(self, inputs): break self.precision = p.dtype if "p" in locals() else torch.float32 if type(inputs) in [list, tuple]: + inputs = [x if x.dtype == torch.long else x.type(self.precision) + for x in inputs] if self.is_cuda: - inputs = [(x.type(self.precision).cuda() if self.is_cuda else - x.type(self.precision)) if x.dtype != torch.long - else x for x in inputs] + inputs = [x.cuda() for x in inputs] return inputs else: if not (inputs.dtype == torch.long): diff --git a/core/NeuralLayers/dropblock.py b/core/NeuralLayers/dropblock.py index ecaa42c..0728c48 100644 --- a/core/NeuralLayers/dropblock.py +++ b/core/NeuralLayers/dropblock.py @@ -20,7 +20,8 @@ class DropBlock(nn.Module): tensor channels. default = False - recommended in paper iterative_p: when True, iteratively increases probability from 0 to p till n_iterations = steps_to_max, and maintains p there after. - steps_to_max: steps to reach p, default = 500000 + default = True + steps_to_max: iterations to reach p, default = 50000 Return: torch.Tensor of shape BCHW @@ -31,8 +32,8 @@ def __init__(self, p: float = 0.1, block_size: int = 5, shared: bool = False, - iterative_p: bool = False, - steps_to_max: int = 500000): + iterative_p: bool = True, + steps_to_max: int = 50000): super(DropBlock, self).__init__() # checks @@ -70,7 +71,7 @@ def __init__(self, "{}".format(type(iterative_p).__name__)) if iterative_p: # steps_to_max = steps to reach p - self.steps_to_max = 500000 + self.steps_to_max = steps_to_max self.register_buffer("n_iterations", torch.Tensor([0]).sum()) def forward(self, tensor): @@ -89,12 +90,14 @@ def forward(self, tensor): pad = self.w//2 if self.shared: c = 1 + mask = torch.ones(n, c, h-2*pad, w-2*pad).to(tensor.device) mask = torch.bernoulli(mask * gamma) mask = F.pad(mask, (pad, pad, pad, pad)) - block_mask = F.conv2d(mask, torch.ones(c, 1, self.w, self.w), - padding=pad, groups=c) + kernel = torch.ones(c, 1, self.w, self.w).to(tensor.device) + block_mask = F.conv2d(mask, kernel, padding=pad, groups=c) block_mask = (block_mask == 0).float().detach() + # norm = count(M)/count_ones(M) norm = block_mask.sum(2, True).sum(3, True) / h / w return tensor * block_mask * norm # A × count(M)/count_ones(M)