Merge pull request #31 from ckkelvinchan/numpy

Change math to numpy

Author: nbei

apps/stylegan_projector.py

@@ -5,10 +5,10 @@
     Ref: https://github.com/rosinality/stylegan2-pytorch/blob/master/projector.py # noqa
 """
 import argparse
-import math
 import os
 
 import mmcv
+import numpy as np
 import torch
 import torch.nn.functional as F
 from mmcv import Config
@@ -124,7 +124,7 @@ def noise_normalize_(noises):
 
 def get_lr(t, initial_lr, rampdown=0.25, rampup=0.05):
     lr_ramp = min(1, (1 - t) / rampdown)
-    lr_ramp = 0.5 - 0.5 * math.cos(lr_ramp * math.pi)
+    lr_ramp = 0.5 - 0.5 * np.cos(lr_ramp * np.pi)
     lr_ramp = lr_ramp * min(1, t / rampup)
     return initial_lr * lr_ramp
 

mmgen/datasets/samplers/distributed_sampler.py

@@ -1,6 +1,6 @@
 from __future__ import division
-import math
 
+import numpy as np
 import torch
 from torch.utils.data import DistributedSampler as _DistributedSampler
 
@@ -25,8 +25,10 @@ def __init__(self,
         self.samples_per_gpu = samples_per_gpu
         # fix the bug of the official implementation
         self.num_samples_per_replica = int(
-            math.ceil(
-                len(self.dataset) * 1.0 / self.num_replicas / samples_per_gpu))
+            int(
+                np.ceil(
+                    len(self.dataset) * 1.0 / self.num_replicas /
+                    samples_per_gpu)))
         self.num_samples = self.num_samples_per_replica * self.samples_per_gpu
         self.total_size = self.num_samples * self.num_replicas
 
@@ -43,9 +45,10 @@ def update_sampler(self, dataset, samples_per_gpu=None):
             self.samples_per_gpu = samples_per_gpu
         # fix the bug of the official implementation
         self.num_samples_per_replica = int(
-            math.ceil(
-                len(self.dataset) * 1.0 / self.num_replicas /
-                self.samples_per_gpu))
+            int(
+                np.ceil(
+                    len(self.dataset) * 1.0 / self.num_replicas /
+                    self.samples_per_gpu)))
         self.num_samples = self.num_samples_per_replica * self.samples_per_gpu
         self.total_size = self.num_samples * self.num_replicas
 

mmgen/datasets/singan_dataset.py

@@ -1,6 +1,5 @@
-import math
-
 import mmcv
+import numpy as np
 import torch
 from torch.utils.data import Dataset
 
@@ -21,33 +20,36 @@ def create_real_pyramid(real, min_size, max_size, scale_factor_init):
         scale_factor_init (float): The initial scale factor.
     """
 
-    num_scales = math.ceil(
-        math.log(
-            math.pow(min_size / min(real.shape[0], real.shape[1]), 1),
-            scale_factor_init)) + 1
+    num_scales = int(
+        np.ceil(
+            np.log(
+                np.power(min_size / min(real.shape[0], real.shape[1]), 1),
+                scale_factor_init))) + 1
 
-    scale2stop = math.ceil(
-        math.log(
-            min([max_size, max([real.shape[0], real.shape[1]])]) /
-            max([real.shape[0], real.shape[1]]), scale_factor_init))
+    scale2stop = int(
+        np.ceil(
+            np.log(
+                min([max_size, max([real.shape[0], real.shape[1]])]) /
+                max([real.shape[0], real.shape[1]]), scale_factor_init)))
 
     stop_scale = num_scales - scale2stop
 
     scale1 = min(max_size / max([real.shape[0], real.shape[1]]), 1)
     real_max = mmcv.imrescale(real, scale1)
-    scale_factor = math.pow(
+    scale_factor = np.power(
         min_size / (min(real_max.shape[0], real_max.shape[1])),
         1 / (stop_scale))
 
-    scale2stop = math.ceil(
-        math.log(
-            min([max_size, max([real.shape[0], real.shape[1]])]) /
-            max([real.shape[0], real.shape[1]]), scale_factor_init))
+    scale2stop = int(
+        np.ceil(
+            np.log(
+                min([max_size, max([real.shape[0], real.shape[1]])]) /
+                max([real.shape[0], real.shape[1]]), scale_factor_init)))
     stop_scale = num_scales - scale2stop
 
     reals = []
     for i in range(stop_scale + 1):
-        scale = math.pow(scale_factor, stop_scale - i)
+        scale = np.power(scale_factor, stop_scale - i)
         curr_real = mmcv.imrescale(real, scale)
         reals.append(curr_real)
 

mmgen/models/architectures/dcgan/generator_discriminator.py

@@ -1,5 +1,4 @@
-import math
-
+import numpy as np
 import torch
 import torch.nn as nn
 from mmcv.cnn import ConvModule, normal_init
@@ -74,7 +73,7 @@ def __init__(self,
         self.noise_size = noise_size
 
         # the number of times for upsampling
-        self.num_upsamples = int(math.log2(output_scale // input_scale))
+        self.num_upsamples = int(np.log2(output_scale // input_scale))
 
         # output 4x4 feature map
         self.noise2feat = ConvModule(
@@ -244,7 +243,7 @@ def __init__(self,
         self.base_channels = base_channels
 
         # the number of times for downsampling
-        self.num_downsamples = int(math.log2(input_scale // output_scale))
+        self.num_downsamples = int(np.log2(input_scale // output_scale))
 
         # build up downsampling backbone (excluding the output layer)
         downsamples = []

mmgen/models/architectures/lsgan/generator_discriminator.py

@@ -1,5 +1,4 @@
-import math
-
+import numpy as np
 import torch
 import torch.nn as nn
 from mmcv.cnn import ConvModule
@@ -74,7 +73,7 @@ def __init__(self,
                 'act', build_activation_layer(default_act_cfg))
 
         # the number of times for upsampling
-        self.num_upsamples = int(math.log2(output_scale // input_scale)) - 2
+        self.num_upsamples = int(np.log2(output_scale // input_scale)) - 2
 
         # build up convolution backbone (excluding the output layer)
         self.conv_blocks = nn.ModuleList()
@@ -254,7 +253,7 @@ def __init__(self,
                 act_cfg=default_act_cfg))
 
         # the number of times for downsampling
-        self.num_downsamples = int(math.log2(input_scale // output_scale)) - 1
+        self.num_downsamples = int(np.log2(input_scale // output_scale)) - 1
 
         # build up downsampling backbone (excluding the output layer)
         curr_channels = base_channels

mmgen/models/architectures/positional_encoding.py

@@ -1,5 +1,4 @@
-import math
-
+import numpy as np
 import torch
 import torch.nn as nn
 
@@ -68,9 +67,9 @@ def get_embedding(num_embeddings,
         # there is a little difference from the original paper.
         half_dim = embedding_dim // 2
         if not div_half_dim:
-            emb = math.log(10000) / (half_dim - 1)
+            emb = np.log(10000) / (half_dim - 1)
         else:
-            emb = math.log(1e4) / half_dim
+            emb = np.log(1e4) / half_dim
         # compute exp(-log10000 / d * i)
         emb = torch.exp(torch.arange(half_dim, dtype=torch.float) * -emb)
         emb = torch.arange(

mmgen/models/architectures/singan/generator_discriminator.py

@@ -1,6 +1,6 @@
-import math
 from functools import partial
 
+import numpy as np
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -65,9 +65,10 @@ def __init__(self,
             F.interpolate, mode='bicubic', align_corners=True)
 
         for scale in range(num_scales + 1):
-            base_ch = min(base_channels * pow(2, math.floor(scale / 4)), 128)
+            base_ch = min(base_channels * pow(2, int(np.floor(scale / 4))),
+                          128)
             min_feat_ch = min(
-                min_feat_channels * pow(2, math.floor(scale / 4)), 128)
+                min_feat_channels * pow(2, int(np.floor(scale / 4))), 128)
 
             self.blocks.append(
                 GeneratorBlock(
@@ -218,9 +219,10 @@ def __init__(self,
         super().__init__()
         self.blocks = nn.ModuleList()
         for scale in range(num_scales + 1):
-            base_ch = min(base_channels * pow(2, math.floor(scale / 4)), 128)
+            base_ch = min(base_channels * pow(2, int(np.floor(scale / 4))),
+                          128)
             min_feat_ch = min(
-                min_feat_channels * pow(2, math.floor(scale / 4)), 128)
+                min_feat_channels * pow(2, int(np.floor(scale / 4))), 128)
             self.blocks.append(
                 DiscriminatorBlock(
                     in_channels=in_channels,

mmgen/models/architectures/singan/positional_encoding.py

@@ -4,10 +4,10 @@
 mentioned in the paper: Positional Encoding as Spatial Inductive Bias in GANs.
 More details can be found in: https://arxiv.org/pdf/2012.05217.pdf
 """
-import math
 from functools import partial
 
 import mmcv
+import numpy as np
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -97,9 +97,10 @@ def __init__(self,
             F.interpolate, mode='bicubic', align_corners=True)
 
         for scale in range(num_scales + 1):
-            base_ch = min(base_channels * pow(2, math.floor(scale / 4)), 128)
+            base_ch = min(base_channels * pow(2, int(np.floor(scale / 4))),
+                          128)
             min_feat_ch = min(
-                min_feat_channels * pow(2, math.floor(scale / 4)), 128)
+                min_feat_channels * pow(2, int(np.floor(scale / 4))), 128)
 
             if scale == 0:
                 in_ch = (

mmgen/models/architectures/stylegan/generator_discriminator_v1.py

@@ -1,7 +1,7 @@
-import math
 import random
 
 import mmcv
+import numpy as np
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -90,7 +90,7 @@ def __init__(self,
         }
 
         # generator backbone (8x8 --> higher resolutions)
-        self.log_size = int(math.log2(self.out_size))
+        self.log_size = int(np.log2(self.out_size))
 
         self.convs = nn.ModuleList()
         self.to_rgbs = nn.ModuleList()
@@ -332,7 +332,7 @@ def forward(self,
             latent = torch.cat([latent, latent2], 1)
 
         curr_log_size = self.log_size if curr_scale < 0 else int(
-            math.log2(curr_scale))
+            np.log2(curr_scale))
         step = curr_log_size - 2
 
         _index = 0
@@ -408,7 +408,7 @@ def __init__(self,
             1024: 16,
         }
 
-        log_size = int(math.log2(in_size))
+        log_size = int(np.log2(in_size))
         self.log_size = log_size
         in_channels = channels[in_size]
 
@@ -495,7 +495,7 @@ def forward(self, input, transition_weight=1., curr_scale=-1):
             torch.Tensor: Predict score for the input image.
         """
         curr_log_size = self.log_size if curr_scale < 0 else int(
-            math.log2(curr_scale))
+            np.log2(curr_scale))
         step = curr_log_size - 2
         for i in range(step, -1, -1):
             index = self.n_layer - i - 1

mmgen/models/architectures/stylegan/generator_discriminator_v2.py

@@ -1,7 +1,7 @@
-import math
 import random
 
 import mmcv
+import numpy as np
 import torch
 import torch.nn as nn
 from mmcv.runner.checkpoint import _load_checkpoint_with_prefix
@@ -152,7 +152,7 @@ def __init__(self,
             fp16_enabled=fp16_enabled)
 
         # generator backbone (8x8 --> higher resolutions)
-        self.log_size = int(math.log2(self.out_size))
+        self.log_size = int(np.log2(self.out_size))
 
         self.convs = nn.ModuleList()
         self.upsamples = nn.ModuleList()
@@ -534,7 +534,7 @@ def __init__(self,
             1024: 16 * channel_multiplier,
         }
 
-        log_size = int(math.log2(in_size))
+        log_size = int(np.log2(in_size))
 
         in_channels = channels[in_size]