text_plain_output_1.png
from PIL import Image

from collections import deque

from torch.utils.data import Dataset, DataLoader

from torch.utils.data.distributed import DistributedSampler

import matplotlib.pyplot as plt

import numpy as np

import os

import pandas as pd

import time

import torch

import torch.nn as nn

import torch.optim as optim

import torch_xla.core.xla_model as xm

import torch_xla.distributed.parallel_loader as pl

import torchvision.models as models

import torchvision.transforms as T

root = '../input/104-flowers-garden-of-eden'
train_df = pd.DataFrame()
folder = os.listdir(root)
for f in folder:
    classes = os.listdir(os.path.join(root, f, 'train'))
    for c in classes:
        images = os.listdir(os.path.join(root, f, 'train', c))
        tmp_df = pd.DataFrame(images, columns=['image_name'])
        tmp_df['class'] = c
        tmp_df['folder'] = f
        tmp_df['type'] = 'train'
        train_df = train_df.append(tmp_df, ignore_index=True)
val_df = pd.DataFrame()
folder = os.listdir(root)
for f in folder:
    classes = os.listdir(os.path.join(root, f, 'val'))
    for c in classes:
        images = os.listdir(os.path.join(root, f, 'val', c))
        tmp_df = pd.DataFrame(images, columns=['image_name'])
        tmp_df['class'] = c
        tmp_df['folder'] = f
        tmp_df['type'] = 'val'
        val_df = val_df.append(tmp_df, ignore_index=True)
classes = list(train_df['class'].unique())
train_df['label'] = train_df['class'].apply(lambda x: classes.index(x))
val_df['label'] = val_df['class'].apply(lambda x: classes.index(x))

test_df = pd.DataFrame()
f = 'jpeg-224x224'
images = os.listdir(os.path.join(root, f, 'test'))
tmp_df = pd.DataFrame(images, columns=['image_name'])
tmp_df['class'] = 'unknown'
tmp_df['folder'] = f
tmp_df['type'] = 'test'
test_df = test_df.append(tmp_df, ignore_index=True)

class flowerDataset(Dataset):
    def __init__(self, df, root='../input/104-flowers-garden-of-eden'):
        self.df = df
        self.root = root
        self.transforms = T.Compose([T.Resize((224, 224)), T.ToTensor()])
    def __getitem__(self, idx):
        img_path = os.path.join(self.root, self.df.iloc[idx]['folder'], self.df.iloc[idx]['type'], self.df.iloc[idx]['class'], self.df.iloc[idx]['image_name'])
        img = Image.open(img_path)
        img_tensor = self.transforms(img)
        target_tensor = torch.tensor(self.df.iloc[idx]['label'], dtype=torch.long)
        return (img_tensor, target_tensor)
    def __len__(self):
        return len(self.df)
class testDataset(Dataset):
    def __init__(self, df, root='../input/104-flowers-garden-of-eden'):
        self.df = df
        self.root = root
        self.transforms = T.Compose([T.Resize((224, 224)), T.ToTensor()])
    def __getitem__(self, idx):
        img_path = os.path.join(self.root, self.df.iloc[idx]['folder'], self.df.iloc[idx]['type'], self.df.iloc[idx]['image_name'])
        img = Image.open(img_path)
        img_tensor = self.transforms(img)
        return (img_tensor, self.df.iloc[idx]['image_name'][:-5])
    def __len__(self):
        return len(self.df)

train_dataset = flowerDataset(train_df)
print(train_dataset.__len__())

train_loader = DataLoader(train_dataset, batch_size = 32, shuffle = True, drop_last = True)
train_iter = iter(train_loader)

images, labels = next(train_iter)
print(images.size())
print(labels.size())

plot_size = 32

fig = plt.figure(figsize=(25, 10))
for idx in np.arange(plot_size):
    ax = fig.add_subplot(4, plot_size/4, idx+1, xticks=[], yticks=[])
    ax.imshow(np.transpose(images[idx], (1, 2, 0)))
    ax.set_title(classes[labels[idx].item()])

def train_net():
    torch.manual_seed(FLAGS['seed'])
    device = xm.xla_device()
    world_size = xm.xrt_world_size()
    train_dataset = flowerDataset(train_df)
    train_sampler = DistributedSampler(train_dataset, num_replicas=world_size, rank=xm.get_ordinal(), shuffle=True)
    train_loader = DataLoader(train_dataset, batch_size=FLAGS['batch_size'], sampler=train_sampler, num_workers=FLAGS['num_workers'], drop_last=True)
    val_dataset = flowerDataset(val_df)
    val_sampler = DistributedSampler(val_dataset, num_replicas=world_size, rank=xm.get_ordinal(), shuffle=True)
    val_loader = DataLoader(val_dataset, batch_size=FLAGS['batch_size'], sampler=val_sampler, num_workers=FLAGS['num_workers'], drop_last=True)
    model = models.resnet18()
    model.load_state_dict(torch.load('/kaggle/input/resnet18/resnet18.pth'))
    model.fc = nn.Linear(512, 104)
    model.to(device)
    optimizer = optim.SGD(model.parameters(), lr=FLAGS['learning_rate'] * world_size, momentum=FLAGS['momentum'], weight_decay=0.0005)
    loss_fn = torch.nn.CrossEntropyLoss()
    def train_loop_fn(loader):
        tracker = xm.RateTracker()
        model.train()
        loss_window = deque(maxlen=FLAGS['log_steps'])
        for x, (data, target) in enumerate(loader):
            optimizer.zero_grad()
            output = model(data)
            loss = loss_fn(output, target)
            loss_window.append(loss.item())
            loss.backward()
            xm.optimizer_step(optimizer)
            tracker.add(FLAGS['batch_size'])
    def val_loop_fn(loader):
        total_samples, correct = (0, 0)
        model.eval()
        for data, target in loader:
            with torch.no_grad():
                output = model(data)
            pred = output.max(1, keepdim=True)[1]
            correct += pred.eq(target.view_as(pred)).sum().item()
            total_samples += data.size()[0]
        accuracy = 100.0 * correct / total_samples
        return accuracy
    for epoch in range(1, FLAGS['num_epochs'] + 1):
        para_loader = pl.ParallelLoader(train_loader, [device])
        train_loop_fn(para_loader.per_device_loader(device))
        para_loader = pl.ParallelLoader(val_loader, [device])
        accuracy = val_loop_fn(para_loader.per_device_loader(device))
        best_accuracy = 0.0
        if accuracy > best_accuracy:
            xm.save(model.state_dict(), 'trained_resnet18_model.pth')
            best_accuracy = accuracy

def _mp_fn(rank, flags):
    global FLAGS
    FLAGS = flags
    torch.set_default_tensor_type('torch.FloatTensor')
    train_start = time.time()
    train_net()
    elapsed_train_time = time.time() - train_start

model = models.resnet18()
model.fc = nn.Linear(512, 104)
model.load_state_dict(torch.load('trained_resnet18_model.pth'))
device = xm.xla_device()
model.to(device)
model.eval()

batch_size = 32
test_dataset = testDataset(test_df)
test_loader = DataLoader(test_dataset, batch_size=batch_size)
n = test_dataset.__len__()

predictions = pd.DataFrame()
for x, (images, names) in enumerate(test_loader):
    images = images.to(device)
    with torch.no_grad():
        output = model(images)
    pred = output.max(1)[1].cpu().numpy()
    predictions = predictions.append(pd.DataFrame(data={'id': names, 'label': pred}), ignore_index=True)
    print('\rProcess {} %'.format(round(100 * x * batch_size / n)), end='')
predictions.head()