snt.AbstractModule in dm-sonnet=2.0.0

Question

I've been trying to run code written in 2017, which uses tensorflow=1.3.0 and dm-sonnet=1.14. Many of the packages required have become unavailable and I was drowning in dependency issues so I decided to rewrite the code in a new version of tensorflow.

I have:

tensorflow=2.3.1
tensorflow-addons=0.11.2
tensorflow-estimator=2.3.0
tensorflow-probability=0.11.1
dm-sonnet=2.0.0

Many of the updates are just simple updating tf1 to tf2 but I have endless problems with dm-sonnet versions.

The code has:

class CosineWeights(snt.AbstractModule):

and I get:

AttributeError: module 'sonnet' has no attribute 'AbstractModule'

I suspect this is just a matter of updating the code of sonnet=1.14 to sonnet 2.0 but I can't find a "translation" for it, or how AbstractModule is used now. Which module comes instead, is it enough to just do snt.Module? I would greatly appreciate any help!

Answer 1

Quote From https://github.com/deepmind/sonnet It fixed my problem. Hope this will help~

Building your own module Sonnet strongly encourages users to subclass snt.Module to define their own modules. Let's start by creating a simple Linear layer called MyLinear:

class MyLinear(snt.Module):

  def __init__(self, output_size, name=None):
    super(MyLinear, self).__init__(name=name)
    self.output_size = output_size