<commit_before>local utils = {}

-- a script to simplify trained net by incorporating every SpatialBatchNormalization to SpatialConvolution
-- and BatchNormalization to Linear
local function BNtoConv(net)
  for i,v in ipairs(net.modules) do
    if v.modules then
      BNtoConv(v)
    else
      if (torch.typename(v):find'nn.SpatialBatchNormalization') and net:get(i-1) and
        (torch.typename(net:get(i-1)):find'SpatialConvolution') then
        local conv = net:get(i-1)
        local bn = v
        net:remove(i)
        local no = conv.nOutputPlane
        local conv_w = conv.weight:view(no,-1)
        cutorch.withDevice(conv_w:getDevice(), function()
           if bn.running_var then
              bn.running_std = bn.running_var:add(bn.eps):pow(-0.5)
           end
           if not conv.bias then
              conv.bias = bn.running_mean:clone():zero()
              conv.gradBias = conv.bias:clone()
           end
          conv_w:cmul(bn.running_std:view(no,-1):expandAs(conv_w))
          conv.bias:add(-1,bn.running_mean):cmul(bn.running_std)
          if bn.affine then
            conv.bias:cmul(bn.weight):add(bn.bias)
            conv_w:cmul(bn.weight:view(no,-1):expandAs(conv_w))
          end
          if conv.resetWeightDescriptors then
             conv:resetWeightDescriptors()
             assert(conv.biasDesc)
          end
        end)
      end
    end
  end
end

local function BNToLinear(net)
  for i,v in ipairs(net.modules) do
    if v.modules then
      BNToLinear(v)
    else
      if (torch.typename(v):find'nn.BatchNormalization') and
        (torch.typename(net:get(i-1)):find'Linear') then
        local linear = net:get(i-1)
        local bn = v
        net:remove(i)
        local no = linear.weight:size(1)
        if bn.running_var then
          bn.running_std = bn.running_var:add(x.eps):pow(-0.5)
        end
        linear.weight:cmul(bn.running_std:view(no,1):expandAs(linear.weight))
        linear.bias:add(-1,bn.running_mean):cmul(bn.running_std)
        if bn.affine then
          linear.bias:cmul(bn.weight):add(bn.bias)
          linear.weight:cmul(bn.weight:view(no,1):expandAs(linear.weight))
        end
      end
    end
  end
end

function utils.foldBatchNorm(net)
  -- works in place!
  BNtoConv(net)
  BNtoConv(net)
  BNToLinear(net)
  assert(#net:findModules'nn.SpatialBatchNormalization' == 0)
  assert(#net:findModules'nn.BatchNormalization' == 0)
  assert(#net:findModules'cudnn.SpatialBatchNormalization' == 0)
  assert(#net:findModules'cudnn.BatchNormalization' == 0)<commit_msg>Extended to volumetric batchnorm<commit_after>local utils = {}

-- a script to simplify trained net by incorporating every Spatial/VolumetricBatchNormalization
-- to Spatial/VolumetricConvolution and BatchNormalization to Linear
local function BNtoConv(net)
  for i,v in ipairs(net.modules) do
    if v.modules then
      BNtoConv(v)
    else
      if net:get(i-1) and 
        ((torch.typename(v):find'nn.SpatialBatchNormalization' and 
          torch.typename(net:get(i-1)):find'SpatialConvolution') or
         (torch.typename(v):find'nn.VolumetricBatchNormalization' and 
          torch.typename(net:get(i-1)):find'VolumetricConvolution')) then
        local conv = net:get(i-1)
        local bn = v
        net:remove(i)
        local no = conv.nOutputPlane
        local conv_w = conv.weight:view(no,-1)
        cutorch.withDevice(conv_w:getDevice(), function()
           if bn.running_var then
              bn.running_std = bn.running_var:add(bn.eps):pow(-0.5)
           end
           if not conv.bias then
              conv.bias = bn.running_mean:clone():zero()
              conv.gradBias = conv.bias:clone()
           end
          conv_w:cmul(bn.running_std:view(no,-1):expandAs(conv_w))
          conv.bias:add(-1,bn.running_mean):cmul(bn.running_std)
          if bn.affine then
            conv.bias:cmul(bn.weight):add(bn.bias)
            conv_w:cmul(bn.weight:view(no,-1):expandAs(conv_w))
          end
          if conv.resetWeightDescriptors then
             conv:resetWeightDescriptors()
             assert(conv.biasDesc)
          end
        end)
      end
    end
  end
end

local function BNToLinear(net)
  for i,v in ipairs(net.modules) do
    if v.modules then
      BNToLinear(v)
    else
      if (torch.typename(v):find'nn.BatchNormalization') and
        (torch.typename(net:get(i-1)):find'Linear') then
        local linear = net:get(i-1)
        local bn = v
        net:remove(i)
        local no = linear.weight:size(1)
        if bn.running_var then
          bn.running_std = bn.running_var:add(x.eps):pow(-0.5)
        end
        linear.weight:cmul(bn.running_std:view(no,1):expandAs(linear.weight))
        linear.bias:add(-1,bn.running_mean):cmul(bn.running_std)
        if bn.affine then
          linear.bias:cmul(bn.weight):add(bn.bias)
          linear.weight:cmul(bn.weight:view(no,1):expandAs(linear.weight))
        end
      end
    end
  end
end

function utils.foldBatchNorm(net)
  -- works in place!
  BNtoConv(net)
  BNtoConv(net)
  BNToLinear(net)
  assert(#net:findModules'nn.SpatialBatchNormalization' == 0)
  assert(#net:findModules'nn.VolumetricBatchNormalization' == 0)
  assert(#net:findModules'nn.BatchNormalization' == 0)
  assert(#net:findModules'cudnn.SpatialBatchNormalization' == 0)
  assert(#net:findModules'cudnn.VolumetricBatchNormalization' == 0)  
  assert(#net:findModules'cudnn.BatchNormalization' == 0)
{'fraction_non_alphanumeric': 0.11889011392870268, 'fraction_numerical': 0.008269018743109152, 'mean_word_length': 2.682679296346414, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '587964', 'n_tokens_mistral': 1877, 'n_tokens_neox': 1845, 'n_words': 308}