| | |
| |
|
| | import torch |
| | import torch.nn as nn |
| | import torch.optim as optim |
| | import torch.nn.functional as F |
| | import numpy as np |
| | import torchvision |
| |
|
| | |
| | from torch.utils.data import Dataset, DataLoader |
| | import matplotlib.pyplot as plt |
| | import torchvision.models as models |
| | import torchvision.transforms as transforms |
| | import torchvision.datasets as datasets |
| |
|
| | import time |
| | import copy |
| | import os |
| |
|
| |
|
| | batch_size = 128 |
| | learning_rate = 1e-3 |
| |
|
| |
|
| | transforms = transforms.Compose([transforms.ToTensor()]) |
| |
|
| |
|
| | train_dataset = datasets.ImageFolder( |
| | root="/input/fruits-360-dataset/fruits-360/Training", transform=transforms |
| | ) |
| |
|
| | test_dataset = datasets.ImageFolder( |
| | root="/input/fruits-360-dataset/fruits-360/Test", transform=transforms |
| | ) |
| |
|
| |
|
| | train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True) |
| |
|
| | test_dataloader = DataLoader(test_dataset, batch_size=batch_size, shuffle=True) |
| |
|
| | device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") |
| |
|
| |
|
| | def imshow(inp, title=None): |
| |
|
| | inp = inp.cpu() if device else inp |
| | inp = inp.numpy().transpose((1, 2, 0)) |
| | mean = np.array([0.485, 0.456, 0.406]) |
| | std = np.array([0.229, 0.224, 0.225]) |
| | inp = std * inp + mean |
| | inp = np.clip(inp, 0, 1) |
| | plt.imshow(inp) |
| |
|
| | if title is not None: |
| | plt.title(title) |
| | plt.pause(0.001) |
| |
|
| |
|
| | images, labels = next(iter(train_dataloader)) |
| | print("images-size:", images.shape) |
| |
|
| | out = torchvision.utils.make_grid(images) |
| | print("out-size:", out.shape) |
| |
|
| |
|
| | imshow(out, title=[train_dataset.classes[x] for x in labels]) |
| |
|
| |
|
| | net = models.resnet18(pretrained=True) |
| |
|
| | net = net.cuda() if device else net |
| |
|
| | net |
| |
|
| | criterion = nn.CrossEntropyLoss() |
| |
|
| | optimizer = optim.SGD(net.parameters(), lr=0.0001, momentum=0.9) |
| |
|
| |
|
| | def accuracy(out, labels): |
| | _, pred = torch.max(out, dim=1) |
| | return torch.sum(pred == labels).item() |
| |
|
| |
|
| | num_ftrs = net.fc.in_features |
| | net.fc = nn.Linear(num_ftrs, 128) |
| | net.fc = net.fc.cuda() if use_cuda else net.fc |
| |
|
| |
|
| | |
| | _epochs = 5 |
| | print_every = 10 |
| | valid_loss_min = np.Inf |
| | val_loss = [] |
| | val_acc = [] |
| | train_loss = [] |
| | train_acc = [] |
| | total_step = len(train_dataloader) |
| |
|
| | for epoch in range(1, n_epochs + 1): |
| | running_loss = 0.0 |
| | correct = 0 |
| | total = 0 |
| |
|
| | print(f"Epoch {epoch}\n") |
| |
|
| | for batch_idx, (data_, target_) in enumerate(train_dataloader): |
| | data_, target_ = data_.to(device), target_.to(device) |
| | optimizer.zero_grad() |
| | outputs = net(data_) |
| | loss = criterion(outputs, target_) |
| | loss.backward() |
| | optimizer.step() |
| |
|
| | running_loss += loss.item() |
| | _, pred = torch.max(outputs, dim=1) |
| | correct += torch.sum(pred == target_).item() |
| | total += target_.size(0) |
| |
|
| | if (batch_idx) % 20 == 0: |
| | print( |
| | "Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}".format( |
| | epoch, n_epochs, batch_idx, total_step, loss.item() |
| | ) |
| | ) |
| |
|
| | train_acc.append(100 * correct / total) |
| | train_loss.append(running_loss / total_step) |
| | print( |
| | f"\ntrain-loss: {np.mean(train_loss):.4f}, train-acc: {(100 * correct/total):.4f}" |
| | ) |
| |
|
| | batch_loss = 0 |
| | total_t = 0 |
| | correct_t = 0 |
| |
|
| | with torch.no_grad(): |
| | net.eval() |
| | for data_t, target_t in test_dataloader: |
| | data_t, target_t = data_t.to(device), target_t.to(device) |
| | outputs_t = net(data_t) |
| | loss_t = criterion(outputs_t, target_t) |
| | batch_loss += loss_t.item() |
| | _, pred_t = torch.max(outputs_t, dim=1) |
| | correct_t += torch.sum(pred_t == target_t).item() |
| | total_t += target_t.size(0) |
| |
|
| | val_acc.append(100 * correct_t / total_t) |
| | val_loss.append(batch_loss / len(test_dataloader)) |
| |
|
| | network_learned = batch_loss < valid_loss_min |
| | print( |
| | f"validation loss: {np.mean(val_loss):.4f}, validation acc: {(100 * correct_t/total_t):.4f}\n" |
| | ) |
| |
|
| | if network_learned: |
| | valid_loss_min = batch_loss |
| | torch.save(net.state_dict(), "resnet.pt") |
| | print("Improvement-Detected, save-model") |
| |
|
| | net.train() |
| |
|
