【Python】Pytorch 自前画像で転移学習 ResNet50

参考にさせていただいたページ
【PyTorch】torchvision.modelsでResNetやEfficientNetの読み込みと分類クラス数の変更、ファインチューニングへの活用 | 日々、学ぶ

全コード

# moduleのimport
import os
import torch
import torch.nn as nn
import torchvision
import matplotlib.pyplot as plt

# work spaceフォルダを指定
ws = ”work spaceフォルダのpath”
os.chdir(ws)

# data_transformの作成
data_transforms = {
    'train': torchvision.transforms.Compose([
        #torchvision.transforms.RandomResizedCrop(224),
        torchvision.transforms.ToTensor(),
        torchvision.transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
    'valid': torchvision.transforms.Compose([
        #torchvision.transforms.Resize(256),
        #torchvision.transforms.CenterCrop(224),
        torchvision.transforms.ToTensor(),
        torchvision.transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
}


# 自前データのフォルダの指定
train_data_dir = trainパスを指定
valid_data_dir = validiationパスを指定


# datasetの作成とサイズの確認
image_datasets = {
  'train': torchvision.datasets.ImageFolder(train_data_dir, transform=data_transforms['train']),
  'valid': torchvision.datasets.ImageFolder(valid_data_dir, transform=data_transforms['valid'])
}

dataloaders = {
  'train': torch.utils.data.DataLoader(image_datasets['train'], batch_size=16, shuffle=True),
  'valid': torch.utils.data.DataLoader(image_datasets['valid'], batch_size=16)
}

dataset_sizes = {
    'train': len(image_datasets['train']),
    'valid': len(image_datasets['valid'])
}


#class数の獲得
class_names = image_datasets['train'].classes


# deviceの定義
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')


# CNNの定義を出力層クラス数の書き換え　（出力層の出力数を ImageNet の 1000 からこのデータセットのクラス数である 2 に置き換える。）
net_ft1 = torchvision.models.resnet50(pretrained=True)
net_ft1.fc = nn.Linear(net_ft1.fc.in_features, len(class_names))
net_ft1 = net_ft1.to(device)


# loss function & parameterの定義
criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(net_ft1.parameters(), lr=0.001, momentum=0.9)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=7, gamma=0.1)


# Training

num_epochs = 20
acc_history_ft1 = {'train': [], 'valid': []}

for epoch in range(num_epochs):
    print('Epoch {}/{}'.format(epoch, num_epochs - 1))
    print('-' * 10)

    for phase in ['train', 'valid']:
        if phase == 'train':
            net_ft1.train()
        else:
            net_ft1.eval()

        running_loss = 0.0
        running_corrects = 0

        for inputs, labels in dataloaders[phase]:
            inputs = inputs.to(device)
            labels = labels.to(device)

            optimizer.zero_grad()

            with torch.set_grad_enabled(phase == 'train'):
                outputs = net_ft1(inputs)
                _, preds = torch.max(outputs, 1)
                loss = criterion(outputs, labels)

                if phase == 'train':
                    loss.backward()
                    optimizer.step()

            running_loss += loss.item() * inputs.size(0)
            running_corrects += torch.sum(preds == labels.data)
        if phase == 'train':
            scheduler.step()

        epoch_loss = running_loss / dataset_sizes[phase]
        epoch_acc = running_corrects.double() / dataset_sizes[phase]
        acc_history_ft1[phase].append(epoch_acc)
        print('{} Loss: {:.4f} Acc: {:.4f}'.format(phase, epoch_loss, epoch_acc))
        
        # model save
        torch.save(net_ft1, 'resnet_test.pth')


# plot training process
train_value = torch.tensor(acc_history_ft1['train'], device = 'cpu')
valid_value = torch.tensor(acc_history_ft1['valid'], device = 'cpu')

fig = plt.figure(figsize=(4,4),dpi=200)
ax = fig.add_subplot()
ax.plot(train_value, label='train')
ax.plot(valid_value, label='valid')
ax.legend()
ax.set_ylim(0, 1)
fig.show()