<jupyter_start><jupyter_text>Install fastai2 from github<jupyter_code>!pip install -Uqq fastcore --upgrade
!pip install -Uqq fastai --upgrade 
!pip install -Uqq git+https://github.com/butchland/fastai_xla_extensions
VERSION = "20200707"  #"20200515" @param ["1.5" , "20200325", "nightly"]
!curl https://raw.githubusercontent.com/pytorch/xla/master/contrib/scripts/env-setup.py -o pytorch-xla-env-setup.py
#!TORCH_SHOW_CPP_STACKTRACES=1 python pytorch-xla-env-setup.py --apt-packages libomp5 libopenblas-dev
!python pytorch-xla-env-setup.py  --version $VERSION --apt-packages libomp5 libopenblas-dev
!pip freeze | grep torch 
!pip freeze | grep fastcore
!pip freeze | grep fastai
import fastai_xla_extensions.core
from fastai.vision.all import *
from fastai_xla_extensions.core import *
default_device()
path = untar_data(URLs.MNIST_SAMPLE)
Path.BASE_PATH = path; path.ls()
(path/'train').ls()<jupyter_output><empty_output><jupyter_text>Lenet with convs and F.max_pool2d<jupyter_code>class MyLenet(nn.Module):
    def __init__(self):
        super(MyLenet, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 3)
        self.conv2 = nn.Conv2d(6,16,3)
        self.hiden4 = nn.Linear(400, 2) # 2 outputs instead of 10
    def forward(self, x):
        x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
        x = F.max_pool2d(F.relu(self.conv2(x)), 2)
        x = x.view(-1, self.num_flat_features(x))
        x = self.hiden4(x)
        return x
    
    def num_flat_features(self, x):
        size = x.size()[1:]  # all dimensions except the batch dimension
        num_features = 1
        for s in size:
            num_features *= s
        return num_features<jupyter_output><empty_output><jupyter_text>Lenet with layers<jupyter_code>class Lenet2(nn.Module):
    def __init__(self):
        super(Lenet2, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 3)
        self.conv2 = nn.Conv2d(6, 16, 3)
        self.fc1 = nn.Linear(400, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 2) # Only 2 outputs instead of 10
    def forward(self, x):
        # Max pooling over a (2, 2) window
        x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
        # If the size is a square you can only specify a single number
        x = F.max_pool2d(F.relu(self.conv2(x)), 2)
        x = x.view(-1, self.num_flat_features(x))
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

    def num_flat_features(self, x):
        size = x.size()[1:]  # all dimensions except the batch dimension
        num_features = 1
        for s in size:
            num_features *= s
        return num_features<jupyter_output><empty_output><jupyter_text>Current TPU implementation callbacks  Load in TPU<jupyter_code>dblock = DataBlock(
    splitter = GrandparentSplitter(),
    item_tfms = Resize(28),
    blocks = (ImageBlock, CategoryBlock),
    get_items = get_image_files,
    get_y = parent_label,
    batch_tfms = []
)

dblock.summary(path)<jupyter_output>Setting-up type transforms pipelines
Collecting items from /root/.fastai/data/mnist_sample
Found 14434 items
2 datasets of sizes 12396,2038
Setting up Pipeline: PILBase.create
Setting up Pipeline: parent_label -> Categorize -- {'vocab': None, 'sort': True, 'add_na': False}

Building one sample
  Pipeline: PILBase.create
    starting from
      /root/.fastai/data/mnist_sample/train/3/35491.png
    applying PILBase.create gives
      PILImage mode=RGB size=28x28
  Pipeline: parent_label -> Categorize -- {'vocab': None, 'sort': True, 'add_na': False}
    starting from
      /root/.fastai/data/mnist_sample/train/3/35491.png
    applying parent_label gives
      3
    applying Categorize -- {'vocab': None, 'sort': True, 'add_na': False} gives
      TensorCategory(0)

Final sample: (PILImage mode=RGB size=28x28, TensorCategory(0))


Setting up after_item: Pipeline: Resize -- {'size': (28, 28), 'method': 'crop', 'pad_mode': 'reflection', 'resamples': (2, 0), 'p': 1.0} -> ToTensor
Setting up b[...]<jupyter_text>Use TPU device  Aquire a TPU device<jupyter_code>import torch_xla.core.xla_model as xm


dls_tpu = dblock.dataloaders(path, device=xm.xla_device())
dls_tpu.vocab
xb, yb = dls_tpu.one_batch()
(type(xb), xb.shape, xb.dtype)
(type(yb), yb.shape, yb.dtype)<jupyter_output><empty_output><jupyter_text>First training on TPU<jupyter_code># from fastai.callback.all import *
from fastai.test_utils import *
lenet_tpu = Lenet2()
tpu_learner = Learner(dls_tpu,
                      lenet_tpu,
                      metrics=accuracy, 
                      loss_func=F.cross_entropy,
                      cbs=[])
tpu_learner.summary()
tpu_learner.show_training_loop()
dls_tpu.device<jupyter_output><empty_output><jupyter_text>Call fitWill fail in `self.loss.backward(); `?<jupyter_code>tpu_learner.fit(1)
tpu_learner.lr_find()
%%time
tpu_learner.fit_one_cycle(1, lr=2.0892961401841602e-06)
tpu_learner.validate()
class DeviceCheckCallback(Callback):
    
    def after_pred(self):
        print(f'xb device: {first(self.learn.xb).device}')
        print(f'model device: {one_param(self.learn.model).device}')

tpu_learner.fit_one_cycle(1, cbs=[DeviceCheckCallback()])<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.12813157391470645, 'fraction_numerical': 0.03442340791738382, 'mean_word_length': 3.5086206896551726, 'pattern_counts': {'":': 0, '<': 28, '<?xml version=': 0, '>': 31, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '7', 'dir': 'jupyter-structured-clean-dedup', 'id': '218693', 'license': 'Apache-2.0', 'n_tokens_mistral': 2049, 'n_tokens_neox': 1939, 'n_words': 402, 'path': 'archive_nbs/Basic_lenet_explorationTPU.ipynb', 'repo_name': 'tyoc213/fastai_xla_extensions', 'size': '5229'}