PyTorch-独自データセット（custom dataset）の作り方【初級 深層学習講座】

!pip install timm

!wget https://pjreddie.com/media/files/mnist_train.tar.gz
!wget https://pjreddie.com/media/files/mnist_test.tar.gz

!mkdir dataset
%cd dataset
!tar -xf ../mnist_train.tar.gz 
!tar -xf ../mnist_test.tar.gz
%cd ..

import torch
import torch.nn.functional as F
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
import glob
import multiprocessing

from tqdm.notebook import tqdm
import timm

class CustomDataset(torch.utils.data.Dataset):
  def __init__(self, files, mode):
    self.files = files
    self.mode = mode # mode 0...train, 1...valid, 2...test(without label)

  def __len__(self):
    return len(self.files)

  def __getitem__(self, idx):
    file = self.files[idx]
    image = Image.open(file)
    image = torch.from_numpy(np.array(image).astype(np.float32))
    image = image.unsqueeze(0)

    if self.mode == 0: 
      pass # ここに、データ拡張などの処理を書く
    
    if self.mode == 2:
      label = None
    else:
      label = int(file[-5])


    return image, label

train_file = glob.glob("dataset/train/*")
test_file = glob.glob("dataset/test/*")

train_dataset = CustomDataset(train_file, mode = 0)
test_dataset = CustomDataset(test_file, mode = 1)

img, label = train_dataset[100]
plt.imshow(img[0])
print(label)

EPOCHS = 10
TRAIN_BATCH_SIZE = 32
TEST_BATCH_SIZE = 32
cpus = multiprocessing.cpu_count()

device = "cuda" if torch.cuda.is_available() else "cpu"
print(device)

train_loader = torch.utils.data.DataLoader(train_dataset, batch_size = TRAIN_BATCH_SIZE, shuffle = True, num_workers = cpus)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size = TEST_BATCH_SIZE, shuffle = False, num_workers = cpus)

Iter = iter(train_loader)
imgs, labels = next(Iter)

fix, ax = plt.subplots(4, 8)
for i in range(4):
  for j in range(8):
    ax[i, j].imshow(imgs[i*8+j][0])
    ax[i, j].set_title(int(labels[i*8+j]))
    ax[i, j].axis('off')
plt.show()

model = timm.create_model("resnet18", pretrained = False, num_classes = 10, in_chans = 1)
model = model.to(device)

optimizer = torch.optim.Adam(params=model.parameters(), lr=0.001)

best_loss = 10**18

for epoch in range (EPOCHS) :
  train_loss = 0
  train_cnt = 0
  model.train()
  print(f"EPOCH{epoch}")
  print("[Train]")
  for imgs, labels in tqdm(train_loader) :
    optimizer.zero_grad()
    output = model(imgs.to(device))
    loss = F.cross_entropy(output, labels.to(device))
    loss.backward()
    train_loss += loss.cpu().detach().numpy()
    train_cnt += 1
    optimizer.step()

  valid_loss = 0
  valid_cnt = 0
  model.eval()
  print("[Valid]")
  for imgs, labels in tqdm(test_loader) :
    with torch.no_grad() :
      output = model(imgs.to(device))
    loss = F.cross_entropy(output, labels.to(device))
    valid_loss += loss.cpu().detach().numpy()
    valid_cnt += 1

  if valid_loss < best_loss :
    best_loss = valid_loss
    torch.save(model.state_dict(), "best.pt")

  print("[Score]")
  print(f"train_loss: {train_loss/train_cnt} valid_loss: {valid_loss/valid_cnt}")

y_pred = []
y_true = []

model.load_state_dict(torch.load("best.pt"))
model.eval()

for imgs, labels in tqdm(test_loader) :
  with torch.no_grad() :
    output = model(imgs.to(device))
  y_pred += (torch.argmax(output, axis=1).tolist())
  y_true += labels.tolist()

from sklearn.metrics import confusion_matrix, accuracy_score, precision_score, classification_report

cm = confusion_matrix(y_true, y_pred)
for e in cm :
  print(e)

print(classification_report(y_true, y_pred)) 

class CustomDataset(torch.utils.data.Dataset):
  def __init__(self, 必要な引数):
    # パラメータの設定などを行う
    self.size = データ数
    self.xxxx = xxxxx
        :

  def __len__(self):
    return len(self.size)

  def __getitem__(self, idx):
    # idxに対応した、データを返す
    # ※torchのテンソルにしておくと楽

    return 返すデータ、返すデータ２