Xla: Problem with BatchNorm

[umgupta]

I modified the MNIST example to run in a single file as below. However, I am having problems using batch norm properly.

This code works fine on CPU and on TPU without batchnorm (accuracy >98%). But batchnorm on TPU gives very poor results. Am I missing something here?

import os
import shutil
import time
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
import torch_xla
import torch_xla_py
import torch_xla_py.data_parallel as dp
import torch_xla_py.utils as xu
import torch_xla_py.xla_model as xm
import numpy as np

datadir = '/home/umang/data'
lr = 1e-2
batch_size = 128
num_epochs = 10
momentum = 0.5
log_steps = 10
use_bn = True  # False
devices =  xm.get_xla_supported_devices() 


class MNIST(nn.Module):

    def __init__(self):
        super(MNIST, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        if use_bn:
            self.bn1 = nn.BatchNorm2d(10)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        if use_bn:
            self.bn2 = nn.BatchNorm2d(20)
        self.fc1 = nn.Linear(320, 50)
        self.fc2 = nn.Linear(50, 10)

    def forward(self, x):
        x = F.relu(F.max_pool2d(self.conv1(x), 2))
        if use_bn:
            x = self.bn1(x)
        x = F.relu(F.max_pool2d(self.conv2(x), 2))
        if use_bn:
            x = self.bn2(x)
        x = x.view(-1, 320)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)


def train_loop_fn(model, loader, device, context):
    loss_fn = nn.NLLLoss()
    # NOTE: multiple optimizers are declared
    optimizer = optim.SGD(
        model.parameters(),
        lr=lr, momentum=momentum
    )
    tracker = xm.RateTracker()
    model.train()
    for x, (data, target) in loader:
        optimizer.zero_grad()
        output = model(data)
        loss = loss_fn(output, target)
        loss.backward()
        xm.optimizer_step(optimizer)
        tracker.add(batch_size)
        if x % log_steps == 0:
            print('[{}]({}) Loss={:.5f} Rate={:.2f}'.format(
                device, x, loss.item(), tracker.rate()))


def test_loop_fn(model, loader, device, context):
    total_samples = 0
    correct = 0
    model.eval()
    for x, (data, target) in loader:
        output = model(data)
        pred = output.max(1, keepdim=True)[1]
        correct += pred.eq(target.view_as(pred)).sum().item()
        total_samples += data.size()[0]
    return correct, total_samples


if __name__ == "__main__":
    train_loader = torch.utils.data.DataLoader(
        datasets.MNIST(
            datadir,
            train=True,
            download=True,
            transform=transforms.Compose([
                transforms.ToTensor(),
                transforms.Normalize((0.1307,), (0.3081,))
            ])),
        batch_size=batch_size,
        shuffle=True
    )
    test_loader = torch.utils.data.DataLoader(
        datasets.MNIST(
            datadir,
            train=False,
            transform=transforms.Compose([
                transforms.ToTensor(),
                transforms.Normalize((0.1307,), (0.3081,))
            ])),
        batch_size=batch_size,
        shuffle=True,
    )

    # Scale learning rate to num cores
    lr = lr * max(len(devices), 1)

    # Pass [] as device_ids to run using the PyTorch/CPU engine.
    # TODO: see if initialized model passing is ok too
    print("creating model")
    model_parallel = dp.DataParallel(MNIST, device_ids=devices)

    accuracy = 0.0
    for epoch in range(1, num_epochs + 1):
        print(f"epoch : {epoch}")
        model_parallel(train_loop_fn, train_loader)
        a = model_parallel(test_loop_fn, test_loader)
        correct, total_samples = zip(*a)
        print('Accuracy={:.2f}%'.format(
            100.0 * np.sum(correct) / np.sum(total_samples)))

    # final test on CPU
    model = model_parallel._models[0].cpu()
    model.eval()
    correct = 0
    total_samples = 0
    for (data, target) in test_loader:
        output = model(data)
        pred = output.max(1, keepdim=True)[1]
        correct += pred.eq(target.view_as(pred)).sum().item()
        total_samples += data.size()[0]
    print(f"Final accuracy {correct/ total_samples}")

[see--]

I can't reproduce on docker pull gcr.io/tpu-pytorch/xla:nightly.
output:

Final accuracy 0.9887

I guess you use an older version?

[see--]

Yes, I can repro with docker pull gcr.io/tpu-pytorch/xla:r0.1. Personally, I just use nightly. A lot of these issues are already fixed.

[dlibenzi]

Yes this was fixed after 0.1

On Sat, Aug 3, 2019, 10:26 see-- [e-mail] wrote:

Yes, I can repro with docker pull gcr.io/tpu-pytorch/xla:r0.1.
Personally, I just use nightly. A lot of these issues are already fixed.

—
You are receiving this because you are subscribed to this thread.
Reply to this email directly, view it on GitHub
https://github.com/pytorch/xla/issues/900?email_source=notifications&email_token=ADR5YXCFO7VYEM75CAWQI4LQCU6JRA5CNFSM4IJCMU7KYY3PNVWWK3TUL52HS4DFVREXG43VMVBW63LNMVXHJKTDN5WW2ZLOORPWSZGOD3PJ4JI#issuecomment-517905957,
or mute the thread
https://github.com/notifications/unsubscribe-auth/ADR5YXCGXLQDA4UF6L7TIWLQCU6JRANCNFSM4IJCMU7A
.

[zcain117]

@umgupta if you start using nightly docker or nightly conda env, make sure your TPU software version is also nightly

[umgupta]

ok, I think this gets fixed when I use nightly env. Thanks for the help.

But what is bothering me now is how does test/test_train_mnist.py working and not running into this problem? It uses the same architecture with batch norm and also uses .train and .eval.

[nosound2]

@umgupta if you start using nightly docker or nightly conda env, make sure your TPU software version is also nightly

Hi, can you please advise, how do you select nightly TPU soft version? In my GCP compute engine dashboard I see only pytorch-0.1 version. There is one nightly, but I believe it is for tf.

[jysohn23]

Hi, can you please advise, how do you select nightly TPU soft version? In my GCP compute engine dashboard I see only pytorch-0.1 version. There is one nightly, but I believe it is for tf.

The nightly TPU version you are seeing is both for PyTorch and TF. So please use that one with pytorch-nightly on the Compute VM side.

UPDATE: We have created a pytorch/xla only nightly version for TPU Software Version field: pytorch-nightly. Please use pytorch-nightly for all nightly features for pytorch/xla on TPU side, and not nightly. If you have existing TPUs running nightly, please delete them and recreate one that has TPU version pytorch-nightly. Thanks!

[jysohn23]

(leaving another comment so everyone gets notified on this thread...)

UPDATE: We have created a pytorch/xla only nightly version for TPU Software Version field: pytorch-nightly. Please use pytorch-nightly for all nightly features for pytorch/xla on TPU side, and not nightly. If you have existing TPUs running nightly, please delete them and recreate one that has TPU version pytorch-nightly. Thanks!

[ailzhang]

Closing this issue as it's resolved in pytorch-nightly. Please feel free to reopen if you have followup questions.

[fhwzl]

hi, I have the problem too with batchnorm enable. But I need use the torch-xla:0.1, anyone know the reason? thx