使用Scala的Type Class模式实现神经网络的问题总结（亦shapeless使用小结）

1. 问题背景

首先演示一下问题。

首先定义一个用作初始化神经网络层的Type Class，名为CanAutoInit，具体代码如下：

trait CanAutoInit[-For] {
  def init(foor: For): Unit
}

然后再定义神经网络层的抽象类Layer和模板类LayerLike，代码如下：

sealed trait LayerLike[+Repr <: Layer] {
  def repr: Repr = this.asInstanceOf[Repr]

  def init(implicit op: CanAutoInit[Repr]): Unit = op.init(repr)
}

sealed trait Layer extends LayerLike[Layer] {}

注意这里为了演示问题的清晰性，我们只在LayerLike中实现一个方法init()。在真实的神经网络模型实现中，还需要forward(), backward()等方法。

接下来我们实现两个Layer的实体类，作为演示的例子，分别为PoolingLayer和DropoutLayer，代码如下：

class DropoutLayer() extends Layer with LayerLike[DropoutLayer]

object DropoutLayer {
  implicit val dropoutLayerCanAutoInit = new CanAutoInit[DropoutLayer] {
    override def init(foor: DropoutLayer): Unit = println("Init in Dropout Layer")
  }
}

class PoolingLayer() extends Layer with LayerLike[PoolingLayer]

object PoolingLayer {
  implicit val poolingLayerCanAutoInit = new CanAutoInit[PoolingLayer] {
    override def init(foor: PoolingLayer): Unit = println("Init in Pooling Layer")
  }
}

有两点注意事项：

一是按照Type Class模式的设计规范，隐式变量dropoutLayerCanAutoInit和poolingLayerCanAutoInit分别定义在DropoutLayer和PoolingLayer的伴生对象中；

二是本文为了清晰地演示问题，仅在init()方法中打印“Init in Dropout Layer”之类的字符串，真实实现中的初始化更加复杂。

准备工作进行到这里。接下来我们分别新建一个DropoutLayer和PoolingLayer，并调用其init()方法，看看代码和结果：

val pa = new PoolingLayer
pa.init

val da = new DropoutLayer
da.init

结果如下：

pa: PoolingLayer = PoolingLayer@bab35d4

Init in Pooling Layer

res0: Unit = ()

da: DropoutLayer = DropoutLayer@1e8927e5

Init in Dropout Layer

res1: Unit = ()

可以看到PoolingLayer和DropoutLayer都正常初始化完成，迄今为止没有出现问题。接下来考虑一个场景：假设我们有一个变量list，其类型为List[Layer]，其中存储了一个神经网络的多个Layer，我们想通过for循环对所有Layer进行初始化。代码和结果如下：

val list = List(pa, da)
list.foreach(c => c.init)

结果如下：

Error:(69, 22) could not find implicit value for parameter op: A$A123.this.CanAutoInit[A$A123.this.Layer]

list.foreach(c => c.init)

  ^

Error:(69, 22) not enough arguments for method init: (implicit op: A$A123.this.CanAutoInit[A$A123.this.Layer])Unit.

Unspecified value parameter op.

list.foreach(c => c.init)

  ^

可以看到编译失败。编译器无法找到类型为CanAutoInit[Layer]的隐式变量，因为我们仅定义过类型为CanAutoInit[DropoutLayer]和CanAutoInit[PoolingLayer]的隐式变量。

2. 解决方法

这是个很严重的问题，困扰了我接近一周的时间。我们可以发现，出现这个问题的原因在于，将DropoutLayer和PoolingLayer实例赋给类型为Layer的变量时，丢失了DropoutLayer和PoolingLayer实例中原来的类型信息。我们当然可以通过Pattern Match来很粗暴地处理这个问题：

def init(l: Layer): Unit = l match {
  case ll: PoolingLayer => ll.init
  case ll: DropoutLayer => ll.init
  case _ => throw new UnsupportedOperationException("Unsupported layer for initialization")
}

不过这种方案存在两个问题，致使其在实际中变得完全不可行：

一是通过init(l: Layer)函数对Layer类的init()方法进行了二次包装，极大地损伤了代码易用性；

二是在实际项目中，我们会实现更多的Layer层，如ReluLayer、SoftmaxLayer、ConvLayer等，对每一个Layer都匹配的话重复性代码太多；

这个问题应该是函数式Type Class编程中很容易遇到的一个问题，之前不可能没有解决方案。果然经过两天的Google，发现Scala的shapeless library中提供了解决这个问题的方法：Coproduct。

由于本人对Shapeless也正在学习中，现在的了解很皮毛，这里就不对Coproduct做详细介绍，仅提供以上问题的方法：

首先将Layer层的定义修改为sealed trait，目的是禁止第三方的Layer实现：

sealed trait Layer extends LayerLike[Layer] {}

然后新建CanAutoInit的伴生对象，并在其中加入以下隐式变量和隐式方法：

object CanAutoInit {
  implicit val cnilCanAutoInit: CanAutoInit[CNil] = new CanAutoInit[CNil] {
    override def init(foor: CNil): Unit = println("Init in CNil")
  }

  implicit def coproductCanAutoInit[H, T <: Coproduct]
  (implicit
   hCanAutoInit: CanAutoInit[H],
   tCanAutoInit: CanAutoInit[T]
  ): CanAutoInit[H :+: T] = new CanAutoInit[H :+: T] {
    override def init(foor: H :+: T): Unit = foor match {
      case Inl(h) => hCanAutoInit.init(h)
      case Inr(t) => tCanAutoInit.init(t)
    }
  }

  implicit def genericCanAutoInit[A, C <: Coproduct]
  (implicit
   generic: Generic.Aux[A, C],
   cCanAutoInit: Lazy[CanAutoInit[C]]
  ): CanAutoInit[A] = new CanAutoInit[A] {
    override def init(foor: A): Unit = cCanAutoInit.value.init(generic.to(foor))
  }
}

然后再次运行之前出问题的代码，就可以运行成功了：

list: List[Layer] = List(A$A129$A$A129$PoolingLayer@38d60a48, A$A129$A$A129$DropoutLayer@66668396)

Init in Pooling Layer

Init in Dropout Layer

res2: Unit = ()

注意，关于不同的隐式变量需要放在哪里的问题，如果尚不清楚的话请在Google上搜索Scala implicit resolution来详细了解。

最终的完整Demo代码如下，可以在IntelliJ中创建Scala Worksheet运行：

import shapeless.{:+:, CNil, Coproduct, Generic, Inl, Inr, Lazy}

trait CanAutoInit[-For] {
def init(foor: For): Unit
}

object CanAutoInit {
implicit val cnilCanAutoInit: CanAutoInit[CNil] = new CanAutoInit[CNil] {
override def init(foor: CNil): Unit = println("Init in CNil")
}

implicit def coproductCanAutoInit[H, T <: Coproduct]
(implicit
hCanAutoInit: CanAutoInit[H],
tCanAutoInit: CanAutoInit[T]
): CanAutoInit[H :+: T] = new CanAutoInit[H :+: T] {
override def init(foor: H :+: T): Unit = foor match {
case Inl(h) => hCanAutoInit.init(h)
case Inr(t) => tCanAutoInit.init(t)
}
}

implicit def genericFamilyEncoder[A, C <: Coproduct]
(implicit
generic: Generic.Aux[A, C],
cCanAutoInit: Lazy[CanAutoInit[C]]
): CanAutoInit[A] = new CanAutoInit[A] {
override def init(foor: A): Unit = cCanAutoInit.value.init(generic.to(foor))
}
}

trait LayerLike[+Repr <: Layer] {
def repr: Repr = this.asInstanceOf[Repr]

def init(implicit op: CanAutoInit[Repr]): Unit = op.init(repr)
}

sealed trait Layer extends LayerLike[Layer] {}

class DropoutLayer() extends Layer with LayerLike[DropoutLayer]

object DropoutLayer {
implicit val dropoutLayerCanAutoInit = new CanAutoInit[DropoutLayer] {
override def init(foor: DropoutLayer): Unit = println("Init in Dropout Layer")
}
}

class PoolingLayer() extends Layer with LayerLike[PoolingLayer]

object PoolingLayer {
implicit val poolingLayerCanAutoInit = new CanAutoInit[PoolingLayer] {
override def init(foor: PoolingLayer): Unit = println("Init in Pooling Layer")
}
}

val pa = new PoolingLayer
pa.init

val da = new DropoutLayer
da.init

val list = List(pa, da)
list.foreach(c => c.init)