MXNet (5) Parameter and Block Naming (上)

在 Gluon 中 Parameter 或 Block 其實有自己的名字，筆者目前使用時唯一的好處就是可以透過名字取得特定層數，進而決定要不要計算梯度。

1. collect_params():

先來建立一個簡單的網路架構:

net = nn.Sequential()
net.add(nn.Dense(5, activation='relu'))
net.add(nn.Dense(10))

接者呼叫看看 collect_params():

print(net.collect_params())
print(type(net.collect_params()))

collect_params() 回傳一個 ParameterDict，我們可以透過對應名稱去取得對應層數的資料。

net 是一個 Sequential layer, 所以 net 的 prefix 是 sequential0_。

一層 Dense layer 對應到一組 weight 和 bias。

# print(net.collect_params())
sequential0_ (
  Parameter dense0_weight (shape=(5, 0), dtype=float32)
  Parameter dense0_bias (shape=(5,), dtype=float32)
  Parameter dense1_weight (shape=(10, 0), dtype=float32)
  Parameter dense1_bias (shape=(10,), dtype=float32)
)


# print(type(net.collect_params()))
<class 'mxnet.gluon.parameter.ParameterDict'>

上面 weight 的 shape 分別是 (5, 0) 和 (10, 0)，這是因為 MXNet 的 lazy evaluation。在其運行前無法判斷各個 input channel size。

最後我們來看看如何取出特定層數的值，只要我們 forward 過網路。對應的 key 是 sequential0_ 下的 key 值(e.g. dense0_weight, dense0_bias...)，sequential0_ 是 parent 的 prefix。

x = nd.ones((10, 2))
net.initialize()
net(x)
print(net.collect_params()["dense0_weight"].data())

2. 關於 Block 的命名規則:

一般實例化 layer 時， MXNet 會幫 layer 定義前綴。

d0 = nn.Dense(10)
print(d0.prefix)           # print 前綴
print(d0.collect_params()) # print 前綴下的變數

# print(d0.prefix)
dense0_

# print(d0.collect_params())
dense0_ (
  Parameter dense0_weight (shape=(10, 0), dtype=float32)
  Parameter dense0_bias (shape=(10,), dtype=float32)
)

假設我們新增一個一樣的 dense layer，它的 prifix 並不會重複。

d1 = nn.Dense(10)
print(d1.prefix) # dense1_

prefix 其實是可以自己設定的，在創建 layer 時可以透過 prefix 參數設定。

3. Name scope:

透過 Name scope 我們可以把 parent 的 prefix 加入 children 的 prefix。

class SimpleNet(nn.HybridBlock):

    def __init__(self, **kwargs):
        super(SimpleNet, self).__init__(**kwargs)
        with self.name_scope():
            self.body = nn.HybridSequential()
            self.body.add(
                nn.Conv2D(20, 3, activation="relu"),
                nn.Conv2D(36, 3, activation="relu"),
            )
            self.output = nn.Dense(10)

    def hybrid_forward(self, F, x, *args, **kwargs):
        x = self.body(x)
        return self.output(x)

來看看 net 底下的命名架構:

net = SimpleNet()
net.initialize()
print(net.collect_params())

simplenet0_ (
  Parameter simplenet0_conv0_weight (shape=(20, 0, 3, 3), dtype=)
  Parameter simplenet0_conv0_bias (shape=(20,), dtype=)
  Parameter simplenet0_conv1_weight (shape=(36, 0, 3, 3), dtype=)
  Parameter simplenet0_conv1_bias (shape=(36,), dtype=)
  Parameter simplenet0_dense0_weight (shape=(10, 0), dtype=float32)
  Parameter simplenet0_dense0_bias (shape=(10,), dtype=float32)
)

這時候取參數的時候需要改變 key 值，key 會加入 parent 的 prefix。

x = nd.ones((1, 3, 28, 28))
net(x)
print(net.collect_params()["simplenet0_conv0_weight"].data())

4. 小小結論:

雖然有點小複雜，但來做個總結：

1. MXNet 建立 Block 時會順便幫 Block 命名。
2. 透過 Block 的名字，我們可以定位 Block。
3. Name scope 會將 parent 的 prefix 加入命名的一部分，增加取用時的便利性及可讀性。