How to solve (dependency) error on Google collab in official example for training own DDSP-VST model?

Question

I am completely new to Colab and I have been trying to get the following example to run in Colab: https://g.co/magenta/train-ddsp-vst

The execution (of the unchanged notebook from the link above) works until the point when DDSP is being imported and then fails with an error that looks like incompatible package versions of tensorflow and keras:

Screenshot of execution

The full error message:

---------------------------------------------------------------------------
ImportError                               Traceback (most recent call last)
/usr/local/lib/python3.10/dist-packages/ipyfilechooser/filechooser.py in _on_select_click(self, _b)
    315             if self._callback is not None:
    316                 try:
--> 317                     self._callback(self)
    318                 except TypeError:
    319                     # Support previous behaviour of not passing self

12 frames
<ipython-input-1-90823fe106f1> in run_after_select(chooser)
    308     def run_after_select(chooser):
    309       drive_dir = chooser.selected_path
--> 310       run_training(drive_dir=drive_dir)
    311 
    312     fc = FileChooser(initial_dir)

<ipython-input-1-90823fe106f1> in run_training(drive_dir)
    336   # Import DDSP
    337   # ------------------------------------------------------------------------------
--> 338   from ddsp.colab import colab_utils
    339   globals()['colab_utils'] = colab_utils
    340 

/usr/local/lib/python3.10/dist-packages/ddsp/__init__.py in <module>
     17 # Module imports.
     18 from ddsp import core
---> 19 from ddsp import dags
     20 from ddsp import effects
     21 from ddsp import losses

/usr/local/lib/python3.10/dist-packages/ddsp/dags.py in <module>
     28 import tensorflow.compat.v2 as tf
     29 
---> 30 tfkl = tf.keras.layers
     31 
     32 # Define Types.

/usr/local/lib/python3.10/dist-packages/tensorflow/python/util/lazy_loader.py in __getattr__(self, item)
     56 
     57   def __getattr__(self, item):
---> 58     module = self._load()
     59     return getattr(module, item)
     60 

/usr/local/lib/python3.10/dist-packages/tensorflow/python/util/lazy_loader.py in _load(self)
     39     """Load the module and insert it into the parent's globals."""
     40     # Import the target module and insert it into the parent's namespace
---> 41     module = importlib.import_module(self.__name__)
     42     self._parent_module_globals[self._local_name] = module
     43 

/usr/lib/python3.10/importlib/__init__.py in import_module(name, package)
    124                 break
    125             level += 1
--> 126     return _bootstrap._gcd_import(name[level:], package, level)
    127 
    128 

/usr/local/lib/python3.10/dist-packages/keras/api/__init__.py in <module>
      6 import sys as _sys
      7 
----> 8 from keras.api import keras
      9 from tensorflow.python.util import module_wrapper as _module_wrapper
     10 

/usr/local/lib/python3.10/dist-packages/keras/api/keras/__init__.py in <module>
     11 
     12 from keras import __version__
---> 13 from keras.api.keras import __internal__
     14 from keras.api.keras import activations
     15 from keras.api.keras import applications

/usr/local/lib/python3.10/dist-packages/keras/api/keras/__internal__/__init__.py in <module>
      6 import sys as _sys
      7 
----> 8 from keras.api.keras.__internal__ import layers
      9 from keras.api.keras.__internal__ import legacy
     10 from tensorflow.python.util import module_wrapper as _module_wrapper

/usr/local/lib/python3.10/dist-packages/keras/api/keras/__internal__/layers/__init__.py in <module>
      6 import sys as _sys
      7 
----> 8 from keras.engine.base_layer import BaseRandomLayer
      9 from keras.layers.preprocessing.image_preprocessing import BaseImageAugmentationLayer
     10 from tensorflow.python.util import module_wrapper as _module_wrapper

/usr/local/lib/python3.10/dist-packages/keras/engine/base_layer.py in <module>
     34 from keras import regularizers
     35 from keras.dtensor import lazy_variable
---> 36 from keras.engine import base_layer_utils
     37 from keras.engine import input_spec
     38 from keras.engine import keras_tensor

/usr/local/lib/python3.10/dist-packages/keras/engine/base_layer_utils.py in <module>
     22 
     23 from keras import backend
---> 24 from keras.dtensor import dtensor_api as dtensor
     25 from keras.utils import control_flow_util
     26 from keras.utils import tf_inspect

ImportError: cannot import name 'dtensor_api' from 'keras.dtensor' (/usr/local/lib/python3.10/dist-packages/keras/dtensor/__init__.py)

This error seems to be thrown somewhere down the line with the following import in def run_training():

from ddsp.colab import colab_utils

As this happens with an official example without changes and another user posted very recently with about a similar (the same?) issue, I suspect some packages must have been updated very recently and this now causes the error.

To fix the error, I have tried to first uninstall and then re-install specific versions of tensorflow and/or keras but this causes further compatibility issues and pip throws an error.

Is there a way to fix such issues from within Colab or does this need to be fixed inside ddsp.collab by the maintainers?

Answer 1

I got the same error and looks like I managed to solve it.

Try running these:

pip install keras==2.10.0
pip install tensorflow==2.10.0
pip install tensorflow-probability==0.18.0