I am implementing a new compression algorithm for the weights of a neural network for the Leela Chess project. the weights are roughly
float32s which I want to compress as small as possible. Error tolerance for this application is
2^-17, so lossy compression is clearly the right answer here. All of the weights are between -5 and 5, but 99.995% are in (-.25,.25) and most reasonably closely clumped around zero.
The basic idea with this algorithm is to turn floats into integer multiples of the error tolerance, and then use a utf-8 inspired encoding to represent small values with only 1 byte.
import numpy as np import bz2 def compress(in_path, out_path): with open(in_path, 'rb') as array: net = np.fromfile(in_path, dtype=np.float32) # Quantize net = np.asarray(net * 2**17, np.int32) # Zigzag encode net = (net >> 31) ^ (net << 1) # To variable length result = np.zeros(len(net)*3, dtype=np.uint8) for i in range(3): big = (net >= 128) << 7 result[i::3] = (net % 128) + big net >>= 7 # Delete non-essential indices zeroes = np.where(result == 0) zeroes = zeroes[np.where(zeroes % 3 != 0)] result = np.delete(result, zeroes) with bz2.open(out_path, 'wb') as out: out.write(result.tobytes()) def decompress(in_path, out_path): with bz2.open(in_path, 'rb') as array: result = np.frombuffer(array.read(), dtype=np.uint8) start_inds = np.where(result<128) not_zeroed = np.ones(len(start_inds), dtype=np.bool) # append zeroe so loop doesn't go out of bounds result = np.append(result, np.zeros(4, dtype=np.uint8)) # Get back fixed length from variable length net = np.zeros(len(start_inds), dtype=np.uint32) for i in range(3): change = (result[start_inds] % 128) * not_zeroed net[np.where(not_zeroed)] *= 128 net += change start_inds += 1 not_zeroed &= result[start_inds] >= 128 # Zigzag decode net = (net >> 1) ^ -(net & 1) print(np.mean(net)) # Un-quantize net = np.asarray(net, np.float32) net /= 2**17 with open(out_path, 'wb') as out: out.write(version) out.write(net.tobytes()) compress('diff.hex','diff.bz2') decompress('diff.bz2','round.hex')
The main type of advice I'm looking for is algorithm and performance advice, but ways to make the code readable are always nice.