def train_imagenet():
    print('==> Preparing data..')
    img_dim = get_model_property('img_dim')
    resize_dim = max(img_dim, 256)
    train_loader = make_train_loader(img_dim, batch_size=FLAGS.batch_size, shuffle=10000)
    test_loader = make_val_loader(img_dim, resize_dim, batch_size=FLAGS.test_set_batch_size)
    torch.manual_seed(42)
    server = xp.start_server(FLAGS.profiler_port)
    device = xm.xla_device()
    model = get_model_property('model_fn')().to(device)
    writer = None
    if xm.is_master_ordinal():
        writer = test_utils.get_summary_writer(FLAGS.logdir)
    optimizer = optim.SGD(model.parameters(), lr=FLAGS.lr, momentum=FLAGS.momentum, weight_decay=0.0001)
    num_training_steps_per_epoch = (trainsize // (FLAGS.batch_size * xm.xrt_world_size()))
    lr_scheduler = schedulers.wrap_optimizer_with_scheduler(optimizer, scheduler_type=getattr(FLAGS, 'lr_scheduler_type', None), scheduler_divisor=getattr(FLAGS, 'lr_scheduler_divisor', None), scheduler_divide_every_n_epochs=getattr(FLAGS, 'lr_scheduler_divide_every_n_epochs', None), num_steps_per_epoch=num_training_steps_per_epoch, summary_writer=writer)
    loss_fn = nn.CrossEntropyLoss()

    def train_loop_fn(loader, epoch):
        train_steps = (trainsize // (FLAGS.batch_size * xm.xrt_world_size()))
        tracker = xm.RateTracker()
        total_samples = 0
        rate_list = []
        model.train()
        for (step, (data, target)) in enumerate(loader):
            optimizer.zero_grad()
            output = model(data)
            loss = loss_fn(output, target)
            loss.backward()
            xm.optimizer_step(optimizer)
            tracker.add(FLAGS.batch_size)
            total_samples += data.size()[0]
            if lr_scheduler:
                lr_scheduler.step()
            if ((step % FLAGS.log_steps) == 0):
                xm.add_step_closure(_train_update, args=(device, step, loss, tracker, epoch, writer))
                test_utils.write_to_summary(writer, step, dict_to_write={'Rate_step': tracker.rate()}, write_xla_metrics=False)
            if (step == train_steps):
                break
        reduced_global = xm.mesh_reduce('reduced_global', tracker.global_rate(), np.mean)
        return (total_samples, reduced_global)

    def test_loop_fn(loader, epoch):
        test_steps = (testsize // (FLAGS.test_set_batch_size * xm.xrt_world_size()))
        (total_samples, correct) = (0, 0)
        model.eval()
        for (step, (data, target)) in enumerate(loader):
            output = model(data)
            pred = output.max(1, keepdim=True)[1]
            correct += pred.eq(target.view_as(pred)).sum()
            total_samples += data.size()[0]
            if ((step % FLAGS.log_steps) == 0):
                xm.add_step_closure(test_utils.print_test_update, args=(device, None, epoch, step))
            if (step == test_steps):
                break
        correct_val = correct.item()
        accuracy_replica = ((100.0 * correct_val) / total_samples)
        accuracy = xm.mesh_reduce('test_accuracy', accuracy_replica, np.mean)
        return (accuracy, accuracy_replica, total_samples)
    train_device_loader = pl.MpDeviceLoader(train_loader, device)
    test_device_loader = pl.MpDeviceLoader(test_loader, device)
    (accuracy, max_accuracy) = (0.0, 0.0)
    training_start_time = time.time()
    for epoch in range(1, (FLAGS.num_epochs + 1)):
        xm.master_print('Epoch {} train begin {}'.format(epoch, test_utils.now()))
        replica_epoch_start = time.time()
        (replica_train_samples, reduced_global) = train_loop_fn(train_device_loader, epoch)
        replica_epoch_time = (time.time() - replica_epoch_start)
        avg_epoch_time_mesh = xm.mesh_reduce('epoch_time', replica_epoch_time, np.mean)
        reduced_global = (reduced_global * xm.xrt_world_size())
        xm.master_print('Epoch {} train end {}, Epoch Time={}, Replica Train Samples={}, Reduced GlobalRate={:.2f}'.format(epoch, test_utils.now(), str(datetime.timedelta(seconds=avg_epoch_time_mesh)).split('.')[0], replica_train_samples, reduced_global))
        (accuracy, accuracy_replica, replica_test_samples) = test_loop_fn(test_device_loader, epoch)
        xm.master_print('Epoch {} test end {}, Reduced Accuracy={:.2f}%, Replica Accuracy={:.2f}%, Replica Test Samples={}'.format(epoch, test_utils.now(), accuracy, accuracy_replica, replica_test_samples))
        max_accuracy = max(accuracy, max_accuracy)
        test_utils.write_to_summary(writer, epoch, dict_to_write={'Accuracy/test': accuracy, 'Global Rate': reduced_global}, write_xla_metrics=False)
        if FLAGS.metrics_debug:
            xm.master_print(met.metrics_report())
    test_utils.close_summary_writer(writer)
    total_train_time = (time.time() - training_start_time)
    xm.master_print('Total Train Time: {}'.format(str(datetime.timedelta(seconds=total_train_time)).split('.')[0]))
    xm.master_print('Max Accuracy: {:.2f}%'.format(max_accuracy))
    xm.master_print('Avg. Global Rate: {:.2f} examples per second'.format(reduced_global))
    return max_accuracy
def train_imagenet():
    print('==> Preparing data..')
    img_dim = get_model_property('img_dim')
    resize_dim = max(img_dim, 256)
    train_loader = make_train_loader(img_dim, batch_size=FLAGS.batch_size, shuffle=10000)
    test_loader = make_val_loader(img_dim, resize_dim, batch_size=FLAGS.test_set_batch_size)
    torch.manual_seed(42)
    server = xp.start_server(FLAGS.profiler_port)
    device = xm.xla_device()
    model = get_model_property('model_fn')().to(device)
    writer = None
    if xm.is_master_ordinal():
        writer = test_utils.get_summary_writer(FLAGS.logdir)
    optimizer = optim.SGD(model.parameters(), lr=FLAGS.lr, momentum=FLAGS.momentum, weight_decay=0.0001)
    num_training_steps_per_epoch = (trainsize // (FLAGS.batch_size * xm.xrt_world_size()))
    lr_scheduler = schedulers.wrap_optimizer_with_scheduler(optimizer, scheduler_type=getattr(FLAGS, 'lr_scheduler_type', None), scheduler_divisor=getattr(FLAGS, 'lr_scheduler_divisor', None), scheduler_divide_every_n_epochs=getattr(FLAGS, 'lr_scheduler_divide_every_n_epochs', None), num_steps_per_epoch=num_training_steps_per_epoch, summary_writer=writer)
    loss_fn = nn.CrossEntropyLoss()

    def train_loop_fn(loader, epoch):
        train_steps = (trainsize // (FLAGS.batch_size * xm.xrt_world_size()))
        tracker = xm.RateTracker()
        total_samples = 0
        rate_list = []
        model.train()
        for (step, (data, target)) in enumerate(loader):
            optimizer.zero_grad()
            output = model(data)
            loss = loss_fn(output, target)
            loss.backward()
            xm.optimizer_step(optimizer)
            tracker.add(FLAGS.batch_size)
            total_samples += data.size()[0]
            if lr_scheduler:
                lr_scheduler.step()
            if ((step % FLAGS.log_steps) == 0):
                xm.add_step_closure(_train_update, args=(device, step, loss, tracker, epoch, writer))
                test_utils.write_to_summary(writer, step, dict_to_write={'Rate_step': tracker.rate()}, write_xla_metrics=False)
            if (step == train_steps):
                break
        reduced_global = xm.mesh_reduce('reduced_global', tracker.global_rate(), np.mean)
        return (total_samples, reduced_global)

    def test_loop_fn(loader, epoch):
        test_steps = (testsize // (FLAGS.test_set_batch_size * xm.xrt_world_size()))
        (total_samples, correct) = (0, 0)
        model.eval()
        for (step, (data, target)) in enumerate(loader):
            output = model(data)
            pred = output.max(1, keepdim=True)[1]
            correct += pred.eq(target.view_as(pred)).sum()
            total_samples += data.size()[0]
            if ((step % FLAGS.log_steps) == 0):
                xm.add_step_closure(test_utils.print_test_update, args=(device, None, epoch, step))
            if (step == test_steps):
                break
        correct_val = correct.item()
        accuracy_replica = ((100.0 * correct_val) / total_samples)
        accuracy = xm.mesh_reduce('test_accuracy', accuracy_replica, np.mean)
        return (accuracy, accuracy_replica, total_samples)
    train_device_loader = pl.MpDeviceLoader(train_loader, device)
    test_device_loader = pl.MpDeviceLoader(test_loader, device)
    (accuracy, max_accuracy) = (0.0, 0.0)
    training_start_time = time.time()
    for epoch in range(1, (FLAGS.num_epochs + 1)):
        xm.master_print('Epoch {} train begin {}'.format(epoch, test_utils.now()))
        replica_epoch_start = time.time()
        (replica_train_samples, reduced_global) = train_loop_fn(train_device_loader, epoch)
        replica_epoch_time = (time.time() - replica_epoch_start)
        avg_epoch_time_mesh = xm.mesh_reduce('epoch_time', replica_epoch_time, np.mean)
        reduced_global = (reduced_global * xm.xrt_world_size())
        xm.master_print('Epoch {} train end {}, Epoch Time={}, Replica Train Samples={}, Reduced GlobalRate={:.2f}'.format(epoch, test_utils.now(), str(datetime.timedelta(seconds=avg_epoch_time_mesh)).split('.')[0], replica_train_samples, reduced_global))
        (accuracy, accuracy_replica, replica_test_samples) = test_loop_fn(test_device_loader, epoch)
        xm.master_print('Epoch {} test end {}, Reduced Accuracy={:.2f}%, Replica Accuracy={:.2f}%, Replica Test Samples={}'.format(epoch, test_utils.now(), accuracy, accuracy_replica, replica_test_samples))
        max_accuracy = max(accuracy, max_accuracy)
        test_utils.write_to_summary(writer, epoch, dict_to_write={'Accuracy/test': accuracy, 'Global Rate': reduced_global}, write_xla_metrics=False)
        if FLAGS.metrics_debug:
            xm.master_print(met.metrics_report())
    test_utils.close_summary_writer(writer)
    total_train_time = (time.time() - training_start_time)
    xm.master_print('Total Train Time: {}'.format(str(datetime.timedelta(seconds=total_train_time)).split('.')[0]))
    xm.master_print('Max Accuracy: {:.2f}%'.format(max_accuracy))
    xm.master_print('Avg. Global Rate: {:.2f} examples per second'.format(reduced_global))
    return max_accuracy