CNN in MNIST with PyTorch (PyTorchの基本メモ)

  • PyTorchのチュートリアルとexampleはとても参考になる。0.4から色々変わっているようだし,改めて情報を整理する。
  • まず今回は,MNISTを動かしながら以下の項目についてメモしておく。
    1. MNISTデータのロード方法。可視化方法。
    2. ネットワーク,学習,テストの書き方。
    3. モデルの学習と保存

MNISTデータの利用方法

  • 生のMNISTはバイナリで扱いが面倒。なので,torchvisionのAPIを活用する。(tf.keras, chainerも同様のAPIがあるようだ。)
  • torchvisionはMNISTの他にも,Fasion-MNIST, EMNIST, Cifar-10, STL-10などいくつかのデータセットのダウンロードAPIを用意している。
  • 注意点として,MINST/Fashin-MNIST/EMNISTは同じクラスで実装されているので,ダウンロードで同じディレクトリを指定すると,既存のデータ(/processed/train.ptなど)がすでにあるのでダウンロードをしないことになる。よって,MNIST,EMNIST,Fashion-MNIST毎にrootディレクトリは変更する。transformはダウンロードしてtrain.ptなどを作る段階で適用するわけじゃないので,transformの違いは問題ない。

ダウンロード

  • 下記のコードを実行してMNISTをダウンロードする。rootはデータセットの格納ディレクトリ, downloadをTrueにするとWebから取得する。
  • trainはTrueにするとroot/processed/train.ptからデータを生成し,Falseの場合にはroot/processed/test.pyから生成する。
  • transformはデータに対して共通で適用する前処のの関数を渡す。関数はPILImageを1つ受け取って変換後の画像を返す関数。
from torchvision import datasets, transforms

def main():
    mnist_train = datasets.MNIST(
        root='/home/hoge/PyTorch_MLdata/EMNIST',
        download=True,
        train=True,
        transform=transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.1307,), (0.3081,))
        ]))
    print(mnist)

    mnist_test = datasets.MNIST(
        root='/home/hoge/PyTorch_MLdata/MNIST',
        download=True,  # If already exist, download is ignored.
        train=False,
        transform=transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.1307,), (0.3081,))
        ]))        

if __name__ == '__main__':
    main()        
  • 上記のコードを実行するとrootに指定したディレクトリに下記が生成されている。
- root_dir
  - processed
    - test.pt                  # テスト用データがPyTorchデータ向けに生成されている。
    - training.pt              # 画像用データがPyTorchデータ向けに生成されている。
  - raw
    - t10k-images-idx3-ubyte   # テスト用画像
    - t10k-labels-idx1-ubyte   # テスト用ラベル
    - train-images-idx3-ubyte  # 学習用データ画像
    - train-labels-idx1-ubyte  # 学習用データラベル

データローダー

  • torchにはデータローダというバッチ学習に便利なものがある。単純化以外にも読み込みのマルチスレッド化なども処理してくれるので活用するべき。
  • DataSetはデータの集合。集合ではるが,getitemなどが実装されている。DataLoaderはDataSetを内部に持っていて,バッチ単位で分割してくれるイテレータ
mnist_train_loader = torch.utils.data.DataLoader(mnist_train,
                                                 batch_size=32,
                                                 shuffle=True)
mnist_test_loader = torch.utils.data.DataLoader(mnist_test,
                                                batch_size=32,
                                                shuffle=True)
  • MNISTのようなデータ以外に,自分で入力集合Xsと出力集合Ysを生成した場合,それをTensorDatasetにして渡せばデータローダは簡単に使える。渡す際にはXs, Ysは0次元目のサイズが揃ったtorch.tensorであることが条件。0次元目を軸にgetitemされるからね。 (公式Documentを見た限りだと,別に渡すテンソルは2つに限定されず,n個渡せば,そのn個をローダで良きに取り出してくれるみたいだ。)
import torch.utils.data as Data
# x, y are input and label data (MUST be Pytorch Tensors)
torch_dataset = Data.TensorDataset(x, y)
loader = Data.DataLoader(
    dataset=torch_dataset,
    batch_size=32,
    shuffle=True,
    num_workers=2 # subprocess for loading
)
for epoch in range(n_epoch):
    for step, (batch_x, batch_y) in enumerate(loader):
        # training using batch_x, batch_y
        ...
  • ところで,データセットから直接サブセットを取得したい場合に,datasetはスライスオペレータをサポートしていないので少し工夫がいる。PyTorchのForcumに参考になるコードがある。
# データセットmnist_testから1000個取り出す場合
imgs, labels = zip(*[mnist_test[i] for i in range(1000)])

MNISTデータ可視化

  • あまり本論とは関係ないけれど,MNISTの各データ,ラベルの形式を理解するためにも可視化をしてみる。
  • datasetの返り値は(image:torch.Tensor, label:torch.Tensor)になっている。(というよりも,datasetの生成時にtransform.ToTensorを使ってそうしている)。
  • imageは(1,28,28)のテンソルなので,numpyに変換してmatplotlibのimshowに渡せば描画してくれる。ただし,imshowは(H,W), (H,W,3), (H,W,4)の形式で,各要素は0..1のfloatか0..255の整数でなければならないことに注意する。
image, label = mnist_train[0]  # Tensor(1,28,28), Tensor(1)
image_np = image.numpy()       
plt.imshow(image_np.reshape(28,28))
plt.show()

Network,学習,テストの書き方

  • exampleを見ると,ネットワークをクラスで,train/testはそれぞれ関数として実装している。
  • ネットワークの書き方は色々あるが,exampleでは学習パラメタを含むものをinitに書いて,純粋な関数(relu, poolingなど)はforward層で書いている。(nn.Dropout2dもパラメタ無いはずでは?と思うが。)
  • train, testの書き方は参考になる。deviceはGPU/CPUでコードを共通化する書き方。to(device)で必要だったらGPUテンソル化してくる。
  • train/testで重要な項目の1つが不要なところで勾配計算をさせないこと。テンソルの計算は基本的に計算グラフが構成され,計算のたびに値が保存されてメモリを消費する。それを避けるために,trainの中ではloss.itemでlossの値を抽出している。running_lossを計算するとき名などにテンソルのまま計算してしまうと無駄にメモリを消費するので避ける。また,testのようにそもそも勾配計算が不要な場合にはwith no_grad()でrequire_gradをFalseにしている。
  • また,PyTorchのクロスエントロピー周りは少し注意が必要。functional.cross_entropyはnll_lossとsoftmaxを組合せた関数。モデル側にsoftmaxを入れて,損失関数には単純に負のlog尤度関数(negative-log-likelihood)を使う場合にはnll_lossを使う。
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.conv2_drop = nn.Dropout2d()
        self.fc1 = nn.Linear(320, 50)
        self.fc2 = nn.Linear(50, 10)

    def forward(self, x):
        x = F.relu(F.max_pool2d(self.conv1(x), 2))
        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
        x = x.view(-1, 320)
        x = F.relu(self.fc1(x))
        x = F.dropout(x, training=self.training)
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)

def train(args, model, device, train_loader, optimizer, epoch):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = F.nll_loss(output, target)
        loss.backward()
        optimizer.step()
        if batch_idx % args.log_interval == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                100. * batch_idx / len(train_loader), loss.item()))

def test(args, model, device, test_loader):
    model.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            # sum up batch loss
            test_loss += F.nll_loss(output, target, reduction='sum').item()
            # get the index of the max log-probability
            pred = output.max(1, keepdim=True)[1]
            correct += pred.eq(target.view_as(pred)).sum().item()

    test_loss /= len(test_loader.dataset)
    print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
        test_loss, correct, len(test_loader.dataset),
        100. * correct / len(test_loader.dataset)))

学習とモデルの保存

  • 後は実際に必要なエポック数だけ学習を回して,モデルを保存するだけ。
model = Net().to(device)
optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum)

for epoch in range(1, args.epochs + 1):
    train(args, model, device, train_loader, optimizer, epoch)
    test(args, model, device, test_loader)

# パラメタだけを保存(推奨)
torch.save(model.state_dict(), 'model_param.pkl')
my_model.load_state_dict(torch.load('model_param.pkl'))

# モデル全体をpkl化
# ファイルが大きいこと,GPU/CPU情報を含む依存オブジェクトになってしまうことから非推奨
torch.save(model, 'model.pkl')  # save
my_model = torch.load('model.pkl') # load