def _pickCompressor():
if Options.isExperimental("zstd"):
try:
from zstd import ZSTD_compress # pylint: disable=I0021,import-error
except ImportError:
general.warning(
"Onefile mode cannot compress without 'zstd' module on '%s'." % getOS()
)
else:
return b"Y", lambda data: ZSTD_compress(data, 9, getJobLimit())
else:
return b"X", lambda data: data
retarget(compressed_resnet18)
print('iteration', i)
optim.zero_grad()
y = compressed_resnet18(image)
reconstruction_loss = ((ref - y)**2).mean()
q_loss = quantization_loss(compressed_resnet18)
loss = reconstruction_loss + q_loss
loss.backward(retain_graph=False)
optim.step()
if i % 10:
with torch.no_grad():
model_size_mb = pre_entropy_size(compressed_resnet18).item() / (
8 * 2**20)
pack_size = len(
ZSTD_compress(
compress_weights(compressed_resnet18)
['packs'].numpy().tobytes()))
pack_size_mb = pack_size / (2**20)
print('reconstruction', reconstruction_loss)
print('q_loss', q_loss)
print('size', model_size_mb, 'MB')
print('pack size', pack_size_mb, 'MB')
def compress(data, compression_level=ZSTD_COMPRESSION_LEVEL):
return ZSTD_compress(data, compression_level)