52.3 向Variable / Layer / DataLoader类添加代码

接下来向DeZero的其他类添加支持GPU的代码。本节将向Variable、Layer和DataLoader类添加支持GPU的代码。

首先将Variable类现有的__init__方法和backward方法修改如下。

dezero/core.py

...   
try: import copy array_types  $=$  (np.ndarray,cupy.ndarray) except ImportError:

array_types  $=$  (np.ndarray)   
class Variable: def__init__(self,data,name  $\equiv$  None): if data is not None: if not isinstance(data, array_types): raise TypeError('{}is not supported'.format(type(data))) ... def backward(self, retain_grad=False, create_graph=False): if self.grad is None: xp  $=$ dezero CUDA.get_array_module(self.data) self.grad  $=$  Variable(xp.ones_like(self.data))

__init__方法的参数data要修改为支持copy.ndearray的情况。为此，当copy被成功导入时，执行array_types=(np.ndearray, copy.ndearray)，将array_types动态修改为需要检查的类型。

backward方法修改自动补全梯度(self.grad)的地方。修改后的方法会根据数据的类型(self.data)创建numpy或cupy的多维数组。

此前Variable类在实例变量data中持有NumPy的多维数组。下面我们来实现用于将数据移动到CPU的to_cpu函数，还有将数据移动到GPU的to_gpu函数。代码如下所示。

dezero/core.py

class Variable:   
def to_cpu(self): if self.data is not None: self.data  $=$ dezero.cuda.as_numpy(self.data) def to_gpu(self): if self.data is not None: self.data  $=$ dezero.cuda.as_cupy(self.data)

上面的代码只是使用了as_numpy函数或as_cupy函数。这样就可以将Variable数据从GPU传输到CPU，或者从CPU传输到GPU了。

下一个是Layer类。Layer是保存参数的类。这些参数是继承了Variable的Parameter类。接下来添加将Layer类的参数传输到CPU或GPU的函数。代码如下所示。

dezero/layers.py

class Layer: def to_cpu(self): for param in self.params(): param.to_cpu() def to_gpu(self): for param in self.params(): param.to_gpu()

最后向 DataLoader 类中添加以下阴影部分的代码。

dezero/datalogcers.py

import numpy as np   
fromdezero import CUDA   
class DataLoader: def__init__(self,dataset，batch_size，shuffle=True,gpu=False): self.batch  $\equiv$  dataset self.batch_size  $=$  batch_size self.shuffle  $=$  shuffle self.data_size  $=$  len(batch) self.max_iter  $=$  math.ceil(self.data_size/batch_size) self.gpu  $=$  gpu self.reset()   
def__next__(self): xp  $=$  CUDA.cupy if self.gpu else np x=xp.array([example[0]for example in batch]) t=xp.array([example[1]for example in batch])

self_iteration  $+= 1$  return x, t   
def to_cpu(self): self.cpu  $\equiv$  False   
def to_cpu(self): self.cpu  $\equiv$  True

DataLoader类的作用是从数据集创建小批量数据，这个小批量数据是在__next__方法中创建的。此前创建的小批量数据是NumPy多维数组。这里根据实例变量gpu的标志位在CuPy和NumPy之间进行切换，创建相应的多维数组。

以上就是针对DeZero的3个类Variable、Layer、DataLoader所做的修改。