<jupyter_start><jupyter_text>Vanilla PyTorch training on TinyImageNet dataset This notebook is intended to show that fixing random seeds leads to the same result in both federated and non-federated cases.<jupyter_code>!pip install -r requirements.txt
from pathlib import Path
import os
import shutil
from torch.utils.data import Dataset
from torch.utils.data import DataLoader
from torch import nn
from torch import optim
import torch.nn.functional as F
import torch
import torchvision.transforms as T
import torchvision
import glob
import tqdm
from PIL import Image
import numpy as np<jupyter_output><empty_output><jupyter_text>Download data<jupyter_code>common_data_folder = Path.cwd() / 'data'
zip_file_path = common_data_folder / 'tiny-imagenet-200.zip'
os.makedirs(common_data_folder, exist_ok=True)
os.system(f'wget --no-clobber http://cs231n.stanford.edu/tiny-imagenet-200.zip'
          f' -O {zip_file_path}')
shutil.unpack_archive(str(zip_file_path), str(common_data_folder))
class TinyImageNetDataset(Dataset):
    """TinyImageNet shard dataset class."""

    NUM_IMAGES_PER_CLASS = 500

    def __init__(self, data_folder: Path, data_type='train', transform=None):
        """Initialize TinyImageNetDataset."""
        self.data_type = data_type
        self._common_data_folder = data_folder
        self._data_folder = os.path.join(data_folder, data_type)
        self.labels = {}  # fname - label number mapping
        self.image_paths = sorted(
            glob.iglob(
                os.path.join(self._data_folder, '**', '*.JPEG'),
                recursive=True
            )
        )
        with open(os.path.join(self._common_data_folder, 'wnids.txt'), 'r') as fp:
            self.label_texts = sorted([text.strip() for text in fp.readlines()])
        self.label_text_to_number = {text: i for i, text in enumerate(self.label_texts)}
        self.fill_labels()
        self.transform = transform

    def __len__(self) -> int:
        """Return the len of the shard dataset."""
        return len(self.image_paths)

    def __getitem__(self, index: int):
        """Return an item by the index."""
        file_path = self.image_paths[index]
        sample = self.read_image(file_path)
        if self.transform:
            sample = self.transform(sample)
        label = self.labels[os.path.basename(file_path)]
        return sample, label

    def read_image(self, path: Path):
        """Read the image."""
        img = Image.open(path)
        return img

    def fill_labels(self) -> None:
        """Fill labels."""
        if self.data_type == 'train':
            for label_text, i in self.label_text_to_number.items():
                for cnt in range(self.NUM_IMAGES_PER_CLASS):
                    self.labels[f'{label_text}_{cnt}.JPEG'] = i
        elif self.data_type == 'val':
            with open(os.path.join(self._data_folder, 'val_annotations.txt'), 'r') as fp:
                for line in fp.readlines():
                    terms = line.split('\t')
                    file_name, label_text = terms[0], terms[1]
                    self.labels[file_name] = self.label_text_to_number[label_text]
normalize = T.Normalize(
    mean=[0.485, 0.456, 0.406],
    std=[0.229, 0.224, 0.225]
)

augmentation = T.RandomApply(
    [T.RandomHorizontalFlip(),
     T.RandomRotation(10),
     T.RandomResizedCrop(64)], 
    p=.8
)

training_transform = T.Compose(
    [T.Lambda(lambda x: x.convert("RGB")),
     T.ToTensor(),
     augmentation,
     normalize]
)

valid_transform = T.Compose(
    [T.Lambda(lambda x: x.convert("RGB")),
     T.ToTensor(),
     normalize]
)
def get_train_loader():
    generator=torch.Generator()
    generator.manual_seed(0)
    train_set = TinyImageNetDataset(common_data_folder / 'tiny-imagenet-200', transform=training_transform)
    return DataLoader(train_set, batch_size=64, shuffle=True, generator=generator)

def get_valid_loader():
    valid_set = TinyImageNetDataset(common_data_folder / 'tiny-imagenet-200', data_type='val', transform=valid_transform)
    return DataLoader(valid_set, batch_size=64)<jupyter_output><empty_output><jupyter_text>Describe the model and optimizer<jupyter_code>class Net(nn.Module):
    def __init__(self):
        torch.manual_seed(0)
        super(Net, self).__init__()
        self.model = torchvision.models.mobilenet_v2(pretrained=True)
        self.model.requires_grad_(False)
        self.model.classifier[1] = torch.nn.Linear(in_features=1280, \
                        out_features=200, bias=True)

    def forward(self, x):
        x = self.model.forward(x)
        return x

model = Net()
optimizer = optim.Adam([x for x in model.parameters() if x.requires_grad], lr=1e-4)
loss_fn = F.cross_entropy
def train():
    torch.manual_seed(0)
    device='cpu'
    
    data_loader = tqdm.tqdm(get_train_loader(), desc="train")
    model.train()
    model.to(device)

    losses = []

    for data, target in data_loader:
        data, target = torch.tensor(data).to(device), torch.tensor(
            target).to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = loss_fn(output, target)
        loss.backward()
        optimizer.step()
        losses.append(loss.detach().cpu().numpy())
        
    return {'train_loss': np.mean(losses),}

def validate():
    torch.manual_seed(0)
    device = torch.device('cpu')
    model.eval()
    model.to(device)
    
    data_loader = tqdm.tqdm(get_valid_loader(), desc="validate")
    val_score = 0
    total_samples = 0

    with torch.no_grad():
        for data, target in data_loader:
            samples = target.shape[0]
            total_samples += samples
            data, target = torch.tensor(data).to(device), \
                torch.tensor(target).to(device, dtype=torch.int64)
            output = model(data)
            pred = output.argmax(dim=1,keepdim=True)
            val_score += pred.eq(target).sum().cpu().numpy()
            
    return {'acc': val_score / total_samples,}
for i in range(5):
    if i == 0:
        name, value = next(iter(validate().items()))
        print(f'{name}: {value:f}')
    
    name, value = next(iter(train().items()))
    print(f'{name}: {value:f}')
    
    name, value = next(iter(validate().items()))
    print(f'{name}: {value:f}')<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1256544502617801, 'fraction_numerical': 0.012057750277645565, 'mean_word_length': 2.888957433682912, 'pattern_counts': {'":': 0, '<': 13, '<?xml version=': 0, '>': 15, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '3', 'dir': 'jupyter-structured-clean-dedup', 'id': '181155', 'license': 'Apache-2.0', 'n_tokens_mistral': 2090, 'n_tokens_neox': 2018, 'n_words': 420, 'path': 'openfl-tutorials/interactive_api/PyTorch_TinyImageNet/workspace/non-federated_case.ipynb', 'repo_name': 'yujiag21/openfl', 'size': '6303'}