浮世絵作者予測

Open Review Competition

賞金: 150,000 参加チーム数: 183 終了まで: 30日

浮世絵作者予測のPyTorchによる実装(LB: 0.660)

はじめに

PyTorchを使った学習/推論コード(LB: 0.660)です。

  • google colaboratoryで実行
  • モデルとしてResNet18を使用
  • データ拡張はhorizontal flipのみ

セットアップ

学習データとテストデータをGoogle Driveから読み込みます。 このコードを実行する前に、これらのデータファイルをGoogle Driveの ProbSpace/ukiyoe/ 以下に保存しておいてください。

from google.colab import drive
drive.mount('/content/drive')
Drive already mounted at /content/drive; to attempt to forcibly remount, call drive.mount("/content/drive", force_remount=True).

画像の重複を確認するために imagehash をインストールします。画像の重複確認を行わないなら以下のインストールは不要です。

!pip install imagehash
Requirement already satisfied: imagehash in /usr/local/lib/python3.6/dist-packages (4.0)
Requirement already satisfied: numpy in /usr/local/lib/python3.6/dist-packages (from imagehash) (1.16.5)
Requirement already satisfied: scipy in /usr/local/lib/python3.6/dist-packages (from imagehash) (1.3.1)
Requirement already satisfied: six in /usr/local/lib/python3.6/dist-packages (from imagehash) (1.12.0)
Requirement already satisfied: pillow in /usr/local/lib/python3.6/dist-packages (from imagehash) (4.3.0)
Requirement already satisfied: pywavelets in /usr/local/lib/python3.6/dist-packages (from imagehash) (1.0.3)
Requirement already satisfied: olefile in /usr/local/lib/python3.6/dist-packages (from pillow->imagehash) (0.46)
import google.colab

import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
from sklearn.model_selection import train_test_split

import torch.cuda
import torch.nn as nn
import torch.optim as optim
import torch.autograd as autograd
import torch.utils.data as dataset
import torchvision.models as models

import albumentations
import imagehash
import argparse
import os
import tqdm
import random
import csv

DATA_DIR = 'drive/My Drive/ProbSpace/ukiyoe'
RANDOM_SEED = 2019

random.seed(RANDOM_SEED)
os.environ['PYTHONHASHSEED'] = str(RANDOM_SEED)
np.random.seed(RANDOM_SEED)
torch.manual_seed(RANDOM_SEED)
torch.cuda.manual_seed(RANDOM_SEED)

torch.torch.backends.cudnn.benchmark = True
torch.torch.backends.cudnn.enabled = True

データの読み込み

それぞれのデータをndarrayとして読み込みます。画像のデータ形式は (width, height, channel) となっているようです。

train_labels = np.load(os.path.join(DATA_DIR, 'ukiyoe-train-labels.npz'))['arr_0']
train_images = np.load(os.path.join(DATA_DIR, 'ukiyoe-train-imgs.npz'))['arr_0']
test_images = np.load(os.path.join(DATA_DIR, 'ukiyoe-test-imgs.npz'))['arr_0']

print('train-labels: shape={}, dtype={}'.format(train_labels.shape, train_labels.dtype))
print('train-images: shape={}, dtype={}'.format(train_images.shape, train_images.dtype))
print('test-images: shape={}, dtype={}'.format(test_images.shape, test_images.dtype))
train-labels: shape=(3158,), dtype=int64
train-images: shape=(3158, 224, 224, 3), dtype=uint8
test-images: shape=(397, 224, 224, 3), dtype=uint8

データの確認

学習データに含まれる各ラベルのデータ数を確認。0, 1, 4, 6は多いですが、5, 7, 8, 9は少なめです。

print('train-labels: min={}, max={}'.format(np.min(train_labels), np.max(train_labels)))
axis = plt.figure().add_subplot(1, 1, 1)
axis.set_xlabel('label')
axis.set_ylabel('# of images')
axis.hist(train_labels)
plt.show()
train-labels: min=0, max=9
images = [[] for _ in range(10)]

for image, label in zip(train_images, train_labels):
  images[label].append(image)

figure = plt.figure(figsize=(8, 18))

for i in range(10):
  for j, img in enumerate(images[i][:5]):
    axis = figure.add_subplot(10, 5, i * 5 + j + 1)

    axis.imshow(img)
    axis.tick_params(labelbottom=False, labelleft=False, bottom=False, left=False)
    axis.set_xlabel(f'class={i}')

plt.show()