When you install Tensorflow in the normal way you will often get a warning message like “The TensorFlow library wasn’t compiled to use SSE4.2 instructions, but these are available on your machine and could speed up CPU computations” or something like that.
This tells you that the Tensorflow computations could be optimised further but these optimisations aren’t used by default because they are only available on some CPU architectures.
Tensorflow has a guide on how to install from source with these features enabled but the process is tedious and it is not the Nix way.
If you look at the Nix derivation for tensorflow you can see options for enabling SSE4.2, AVX2 and FMA support; the question is “how do I enable these in my own decleration?”
Here is an extract from the old version of my nix file:
pythonPackageList = python-packages: with python-packages; [
pandas
scikitlearn
Keras
tensorflow
numpy
h5py
scipy
jupyter
conda
statsmodels
matplotlib
seaborn
nltk
beautifulsoup4
flask
psycopg2
];
python3packages = pkgs.python3.withPackages pythonPackageList;
All we need to do is change the tensorflow variable to be what we want it to be:
python3PackageOverrides = pkgs.python3Packages.override (oldAttrs:
{ overrides = self: super: {
tensorflow = super.tensorflow.override {
sse42Support = true;
avx2Support = true;
fmaSupport = true;
};
};
});
pythonPackageList = python-packages: with python-packages; [
pandas
scikitlearn
Keras
python3PackageOverrides.tensorflow
numpy
h5py
scipy
jupyter
conda
statsmodels
matplotlib
seaborn
nltk
beautifulsoup4
flask
psycopg2
];
python3packages = pkgs.python3.withPackages pythonPackageList;
Simple! When you know how…
I don’t notice significant speedups with this change, but it does mean I don’t get the warning message any more. Disabling the warning messages would probably have been easier.