#!/usr/bin/env python3
#-*- coding:utf-8 -*-

import argparse
# import logging
import os
import cv2
import numpy as np
import torch
from pfld.utils import plot_pose_cube
from torch.utils.data import DataLoader
from torchvision import transforms
from tensorboardX import SummaryWriter

from data.datasets import WLFWDatasets
from pfld.pfld import PFLDInference
from pfld.loss import PFLDLoss, MSELoss, SmoothL1, WingLoss
from pfld.utils import AverageMeter

def print_args(args):
    for arg in vars(args):
        s = arg + ': ' + str(getattr(args, arg))
        print(s)

def save_checkpoint(state, filename='checkpoint.pth.tar'):
    torch.save(state, filename)
    print('Save checkpoint to {0:}'.format(filename))

def train(train_loader, plfd_backbone, criterion, optimizer,
          epoch):
    losses = AverageMeter()
    plfd_backbone.train(True)
    for img, landmark_gt, attribute_gt, euler_angle_gt,type_flag in train_loader:
        img.requires_grad = False
        img = img.cuda()
        optimizer.zero_grad()
        attribute_gt.requires_grad = False
        attribute_gt = attribute_gt.cuda()

        landmark_gt.requires_grad = False
        landmark_gt = landmark_gt.cuda()

        euler_angle_gt.requires_grad = False
        euler_angle_gt = euler_angle_gt.cuda()

        type_flag.requires_grad = False
        type_flag = type_flag.cuda()

        plfd_backbone = plfd_backbone.cuda()

        pose, landmarks = plfd_backbone(img)
        # attribute_gt, landmark_gt, euler_angle_gt, angle, landmarks, train_batchsize)
        lds_loss, pose_loss = criterion(attribute_gt, landmark_gt, euler_angle_gt, type_flag,
                                                       pose, landmarks, args.train_batchsize)
        # pose_loss = pose_criterion(pose, euler_angle_gt)
        # lds_loss = pose_criterion(landmarks, landmark_gt)
        total_loss = pose_loss + lds_loss

        total_loss.backward()
        optimizer.step()

        losses.update(total_loss.item())
    return pose_loss, lds_loss

def validate(wlfw_val_dataloader, plfd_backbone, args):
    plfd_backbone.eval()

    with torch.no_grad():
        losses_NME = []
        losses_ION = []
        pose_losses_MAE = []
        for img, landmark_gt, attribute_gt, euler_angle_gt, type_flag  in wlfw_val_dataloader:
            img.requires_grad = False
            img = img.cuda(non_blocking=True)

            attribute_gt.requires_grad = False
            attribute_gt = attribute_gt.cuda(non_blocking=True)

            landmark_gt.requires_grad = False
            landmark_gt = landmark_gt.cuda(non_blocking=True)

            euler_angle_gt.requires_grad = False
            euler_angle_gt = euler_angle_gt.cuda(non_blocking=True)

            type_flag.requires_grad = False
            type_flag = type_flag.cuda()

            plfd_backbone = plfd_backbone.cuda()

            pose, landmarks = plfd_backbone(img)

            if args.check:
                landms_tmp = landmarks.cpu().numpy()
                landms_tmp = landms_tmp.reshape(landms_tmp.shape[0], -1, 2)
                for i in range(landms_tmp.shape[0]):
                    # show result
                    show_img = np.array(np.transpose(img[i].cpu().numpy(), (1, 2, 0)))
                    show_img = (show_img * 255).astype(np.uint8)
                    np.clip(show_img, 0, 255)
                    draw = show_img.copy()
                    yaw = pose[i][0] * 180 / np.pi
                    pitch = pose[i][1] * 180 / np.pi
                    roll = pose[i][2] * 180 / np.pi
                    pre_landmark = landms_tmp[i] * [112, 112]
                    for (x, y) in pre_landmark.astype(np.int8):
                        cv2.circle(draw, (int(x), int(y)), 1, (0, 255, 0), 1)
                    draw = plot_pose_cube(draw, yaw, pitch, roll, size=draw.shape[0] // 2)

                    cv2.imshow("check", draw)
                    cv2.waitKey(0)

            landms_const = torch.tensor(-2).cuda()
            pose_const = torch.tensor(-1).cuda()

            pos1 = type_flag == landms_const

            landmarks = landmarks[pos1]
            landmark_gt = landmark_gt[pos1]
            if landmarks.shape[0] > 0:
                loss = torch.mean(
                    torch.sqrt(torch.sum((landmark_gt * 112 - landmarks  * 112)**2, dim=1))
                    )

                landmarks = landmarks.cpu().numpy()
                landmarks = landmarks.reshape(landmarks.shape[0], -1, 2)
                landmark_gt = landmark_gt.reshape(landmark_gt.shape[0], -1, 2).cpu().numpy()

                error_diff = np.sum(np.sqrt(np.sum((landmark_gt - landmarks) ** 2, axis=2)), axis=1)
                interocular_distance = np.sqrt(np.sum((landmarks[:, 60, :] - landmarks[:,72, :]) ** 2, axis=1))
                # interpupil_distance = np.sqrt(np.sum((landmarks[:, 60, :] - landmarks[:, 72, :]) ** 2, axis=1))
                error_norm = np.mean(error_diff / interocular_distance)

                losses_NME.append(loss.cpu().numpy())
                losses_ION.append(error_norm)

            pos2 = type_flag == pose_const
            pose_p = pose[pos2]
            pose_t = euler_angle_gt[pos2]
            if pose_p.shape[0] > 0:
                pose_loss = torch.mean(abs(pose_t - pose_p) * 180 / np.pi)
                pose_losses_MAE.append(pose_loss.cpu().numpy())

        return np.mean(pose_losses_MAE), np.mean(losses_NME), np.mean(losses_ION)

def adjust_learning_rate(optimizer, initial_lr, step_index):

    lr = initial_lr * (0.1 ** (step_index))
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr
    return lr

def main(args):
    print_args(args)

    plfd_backbone = PFLDInference().cuda()
    if args.resume:
        try:
            plfd_backbone.load_state_dict(torch.load(args.resume, map_location=lambda storage, loc: storage))
            print("load %s successfully ! "%args.resume)
        except KeyError:
            plfd_backbone = torch.nn.DataParallel(plfd_backbone)
            plfd_backbone.load_state_dict(torch.load(args.resume))

    print(plfd_backbone)
    step_epoch = [int(x) for x in args.step.split(',')]

    if args.loss == 'mse':
        criterion = MSELoss()
    elif args.loss == 'sommthl1':
        criterion = SmoothL1()
    elif args.loss == 'wing':
        criterion = WingLoss()
    elif args.loss == 'pfld':
        criterion = PFLDLoss()
    cur_lr = args.base_lr
    optimizer = torch.optim.Adam(
        plfd_backbone.parameters(),
        lr=args.base_lr,
        weight_decay=args.weight_decay)

    # SGD optimizer is very sensitive to the magnitude of loss,
    # which is likely to be non convergent or nan, I recommend Adam.
    # optimizer = torch.optim.SGD(plfd_backbone.parameters(), lr=args.base_lr, momentum=0.9, weight_decay=args.weight_decay)

    train_transform = transforms.Compose([
        transforms.RandomGrayscale(p=0.2),
        transforms.ColorJitter(0.2, 0.2, 0.2, 0.2),
        transforms.ToTensor(),
        transforms.RandomErasing(p=0.2, scale=(0.02, 0.33), ratio=(0.3, 3.3), value=(0.5, 0.5, 0.5)),
    ])
    wlfwdataset = WLFWDatasets(args.dataroot, train_transform)
    dataloader = DataLoader(
        wlfwdataset,
        batch_size=args.train_batchsize,
        shuffle=True,
        num_workers=args.workers,
        drop_last=False)
    val_transform = transforms.Compose([transforms.ToTensor()])
    wlfw_val_dataset = WLFWDatasets(args.val_dataroot, val_transform)
    wlfw_val_dataloader = DataLoader(
        wlfw_val_dataset,
        batch_size=args.val_batchsize,
        shuffle=False,
        num_workers=args.workers)

    step_index = 0
    val_pose_loss, val_lds_nme_loss, val_lds_ion_loss = validate(wlfw_val_dataloader, plfd_backbone, args)
    print('Epoch: %d, val pose MAE:%6.4f, val lds NME:%6.4f, val lds ION:%6.4f, lr:%8.6f' % (step_index, val_pose_loss, val_lds_nme_loss, val_lds_ion_loss, cur_lr))

    writer = SummaryWriter(args.tensorboard)
    for epoch in range(args.start_epoch, args.end_epoch + 1):
        train_pose_loss, train_lds_loss = train(dataloader, plfd_backbone,
                                      criterion, optimizer, epoch)
        filename = os.path.join(
            str(args.snapshot), "checkpoint_epoch_" + str(epoch) + '.pth')
        save_checkpoint(plfd_backbone.state_dict(),filename)
        val_pose_loss, val_lds_nme_loss, val_lds_ion_loss = validate(wlfw_val_dataloader, plfd_backbone, args)
        if epoch in step_epoch:
            step_index += 1
            cur_lr = adjust_learning_rate(optimizer, args.base_lr, step_index)

        print('Epoch: %d, train pose loss: %6.4f, train lds loss:%6.4f, val pose MAE:%6.4f, val lds NME:%6.4f, val lds ION:%6.4f, lr:%8.6f'%(epoch, train_pose_loss, train_lds_loss, val_pose_loss, val_lds_nme_loss, val_lds_ion_loss,cur_lr))
        writer.add_scalar('data/pose_loss', train_pose_loss, epoch)
        writer.add_scalars('data/loss', {'val pose loss': val_pose_loss, 'val lds nme loss': val_lds_nme_loss, 'val lds ion loss': val_lds_ion_loss, 'train loss': train_lds_loss}, epoch)

    writer.close()

def parse_args():
    parser = argparse.ArgumentParser(description='Trainging Template')
    # general
    parser.add_argument('-j', '--workers', default=8, type=int)

    # training
    ##  -- optimizer
    parser.add_argument('--base_lr', default=0.01, type=float)
    parser.add_argument('--weight-decay', '--wd', default=1e-4, type=float)
    parser.add_argument('--step', default="30,60", help="lr decay", type=str)

    # -- epoch
    parser.add_argument('--start_epoch', default=1, type=int)
    parser.add_argument('--end_epoch', default=80, type=int)

    #loss
    parser.add_argument('--loss', default="wing",help="strongly recommend wing loss, other loss function has been abandoned", type=str)
    # -- snapshot、tensorboard log and checkpoint
    parser.add_argument('--snapshot',default='./models/checkpoint/snapshot/',type=str,metavar='PATH')
    parser.add_argument('--tensorboard', default="./models/checkpoint/tensorboard", type=str)
    parser.add_argument('--resume', default='./models/checkpoint/model-wing/checkpoint_epoch_5.pth', type=str, metavar='PATH')  # TBD

    # --dataset
    parser.add_argument('--dataroot',default='/home/unaguo/hanson/data/landmark/WFLW191104/train_data/list.txt',type=str,metavar='PATH')
    parser.add_argument('--val_dataroot',default='/home/unaguo/hanson/data/landmark/WFLW191104/test_data/list.txt',type=str,metavar='PATH')
    parser.add_argument('--train_batchsize', default=128, type=int)
    parser.add_argument('--val_batchsize', default=8, type=int)
    parser.add_argument('--check', default=False, type=bool)
    args = parser.parse_args()
    return args

if __name__ == "__main__":
    args = parse_args()
    main(args)


# [AI训练营]paddleclas实现图像分类baseline

## 0 项目背景
> 对不同的场景进行识别和分类

PS:

本项目属于本次课程的大作业的一部分，希望大家可以学会使用paddleclas实现图像分类。

## 本大作业任务

1、选择一个心仪的数据集
2、运行项目，能够跑通项目即可达到结业要求
3、记得生成新版本，公开项目哦~

**加分项：**
4、更换网络、进行数据预处理和调参

## 1 项目实现简介
正如标题，采用paddleclas套件实现分类[30分钟玩转PaddleClas（尝鲜版）](https://gitee.com/paddlepaddle/PaddleClas/blob/release/2.2/docs/zh_CN/tutorials/quick_start_new_user.md)
查看套件，可以知道
实现分类，仅仅需要我们将数据集提取为如下这种格式的txt文件即可（当然我们并不需要更大的训练集）
![](https://ai-studio-static-online.cdn.bcebos.com/012e5e2e7a204f77ab2ad9759860978d6b19d55946dd4ee2adea6a4e9736d121)

PS：
如有需要参考项目可看：
[基于PaddleClas2.2的从零到落地安卓部署的奥特曼分类实战](https://aistudio.baidu.com/aistudio/projectdetail/2219455)
[iFLYTEK基于PaddleClas2.2的广告分类baseline非官方](https://aistudio.baidu.com/aistudio/projectdetail/2235340)

## 2 数据集介绍

本次数据集有五个可供大家选择。分别是：
1. 猫12分类
2. 垃圾40分类
3. 场景5分类
4. 食物5分类
5. 蝴蝶20分类
我这里使用的是官方给的场景分类
**数据集：都是不同类别的文件夹下放置了对应文件夹名字的类别图片**

```python
# 先导入库
from sklearn.utils import shuffle
import os
import pandas as pd
import numpy as np
from PIL import Image
import paddle
import paddle.nn as nn
import random
```

```python
# 忽略（垃圾）警告信息
# 在python中运行代码经常会遇到的情况是——代码可以正常运行但是会提示警告，有时特别讨厌。
# 那么如何来控制警告输出呢？其实很简单，python通过调用warnings模块中定义的warn()函数来发出警告。我们可以通过警告过滤器进行控制是否发出警告消息。
import warnings
warnings.filterwarnings("ignore")
```

### 2.1 解压数据集，查看数据的结构

```python
# 项目挂载的数据集先解压出来，待解压完毕，刷新后可发现左侧文件夹根目录出现五个zip
!unzip -oq /home/aistudio/data/data103736/五种图像分类数据集.zip
```

左侧可以看到如图所示五个zip
![](https://ai-studio-static-online.cdn.bcebos.com/f8bc5b21a0ba49b4b78b6e7b18ac0341dfb14cf545b14c83b1f597b6ee8109bb)

```python
# 本项目以食物分类为例进行介绍，因为分类大多数情况下是不存在标签文件的，猫分类已经有了标签，省去了数据处理的操作
# (此处需要你根据自己的选择进行解压对应的文件)
# 解压完毕左侧出现文件夹，即为需要分类的文件
!unzip -oq /home/aistudio/场景5分类.zip
```

```python
# 查看结构，正为一个类别下有一系列对应的图片
!tree scenes/
```

    5 directories, 3498 files

**五类食物图片**
1. beef_carpaccio
2. baby_back_ribs
3. beef_tartare
4. apple_pie
5. baklava

具体结构如下：
```
foods/
├── apple_pie
│   ├── 1005649.jpg
│   ├── 1011328.jpg
│   ├── 101251.jpg
```

### 2.2 拿到总的训练数据txt

```python
import os
# -*- coding: utf-8 -*-
# 根据官方paddleclas的提示，我们需要把图像变为两个txt文件
# train_list.txt（训练集）
# val_list.txt（验证集）
# 先把路径搞定 比如：foods/beef_carpaccio/855780.jpg ,读取到并写入txt

# 根据左侧生成的文件夹名字来写根目录
dirpath = "scenes"
# 先得到总的txt后续再进行划分，因为要划分出验证集，所以要先打乱，因为原本是有序的
def get_all_txt():
    all_list = []
    i = 0 # 标记总文件数量
    j = 0 # 标记文件类别
    for root,dirs,files in os.walk(dirpath): # 分别代表根目录、文件夹、文件
        for file in files:
            i = i + 1
            # 文件中每行格式： 图像相对路径      图像的label_id（数字类别）（注意：中间有空格）。
            imgpath = os.path.join(root,file)
            all_list.append(imgpath+" "+str(j)+"\n")

        j = j + 1

    allstr = ''.join(all_list)
    f = open('all_list.txt','w',encoding='utf-8')
    f.write(allstr)
    return all_list , i
all_list,all_lenth = get_all_txt()
print(all_lenth)
```

    3499

### 2.3 数据打乱

```python
# 把数据打乱
all_list = shuffle(all_list)
allstr = ''.join(all_list)
f = open('all_list.txt','w',encoding='utf-8')
f.write(allstr)
print("打乱成功，并重新写入文本")
```

    打乱成功，并重新写入文本

### 2.4 数据划分

```python
# 按照比例划分数据集 食品的数据有5000张图片，不算大数据，一般9:1即可
train_size = int(all_lenth * 0.9)
train_list = all_list[:train_size]
val_list = all_list[train_size:]

print(len(train_list))
print(len(val_list))
```

    3149
    350

```python
# 运行cell，生成训练集txt
train_txt = ''.join(train_list)
f_train = open('train_list.txt','w',encoding='utf-8')
f_train.write(train_txt)
f_train.close()
print("train_list.txt 生成成功！")

# 运行cell，生成验证集txt
val_txt = ''.join(val_list)
f_val = open('val_list.txt','w',encoding='utf-8')
f_val.write(val_txt)
f_val.close()
print("val_list.txt 生成成功！")
```

    train_list.txt 生成成功！
    val_list.txt 生成成功！

## 3 安装paddleclas

数据集核实完搞定成功的前提下，可以准备更改原文档的参数进行实现自己的图片分类了！

这里采用paddleclas的2.2版本，好用！

```python
# 先把paddleclas安装上再说
# 安装paddleclas以及相关三方包(好像studio自带的已经够用了，无需安装了)
!git clone https://gitee.com/paddlepaddle/PaddleClas.git -b release/2.2
# 我这里安装相关包时，花了30几分钟还有错误提示，不管他即可
#!pip install --upgrade -r PaddleClas/requirements.txt -i https://mirror.baidu.com/pypi/simple
```

    Cloning into 'PaddleClas'...
    remote: Enumerating objects: 538, done.
    remote: Counting objects: 100% (538/538), done.
    remote: Compressing objects: 100% (323/323), done.
    remote: Total 15290 (delta 344), reused 349 (delta 210), pack-reused 14752
    Receiving objects: 100% (15290/15290), 113.56 MiB | 10.68 MiB/s, done.
    Resolving deltas: 100% (10236/10236), done.
    Checking connectivity... done.

```python
#因为后续paddleclas的命令需要在PaddleClas目录下，所以进入PaddleClas根目录，执行此命令
%cd PaddleClas
!ls
```

    /home/aistudio/PaddleClas
    dataset  hubconf.py   MANIFEST.in    README_ch.md  requirements.txt
    deploy	 __init__.py  paddleclas.py  README_en.md  setup.py
    docs	 LICENSE      ppcls	     README.md	   tools

```python
# 将图片移动到paddleclas下面的数据集里面
# 至于为什么现在移动，也是我的一点小技巧，防止之前移动的话，生成的txt的路径是全路径，反而需要去掉路径的一部分
!mv ../scenes/ dataset/
```

```python
# 挪动文件到对应目录
!mv ../all_list.txt dataset/scenes
!mv ../train_list.txt dataset/scenes
!mv ../val_list.txt dataset/scenes
```

### 3.1 修改配置文件
#### 3.1.1
主要是以下几点：分类数、图片总量、训练和验证的路径、图像尺寸、数据预处理、训练和预测的num_workers: 0

路径如下：
>PaddleClas/ppcls/configs/quick_start/new_user/ShuffleNetV2_x0_25.yaml

（主要的参数已经进行注释，一定要过一遍）

```
# global configs
Global:
  checkpoints: null
  pretrained_model: null
  output_dir: ./output/
  # 使用GPU训练
  device: gpu
  # 每几个轮次保存一次
  save_interval: 1
  eval_during_train: True
  # 每几个轮次验证一次
  eval_interval: 1
  # 训练轮次
  epochs: 20
  print_batch_step: 1
  use_visualdl: True #开启可视化（目前平台不可用）
  # used for static mode and model export
  # 图像大小
  image_shape: [3, 224, 224]
  save_inference_dir: ./inference
  # training model under @to_static
  to_static: False

# model architecture
Arch:
  # 采用的网络
  name: ResNet50
  # 类别数 多了个0类 0-5 0无用
  class_num: 6

# loss function config for traing/eval process
Loss:
  Train:

    - CELoss:
        weight: 1.0
  Eval:
    - CELoss:
        weight: 1.0

Optimizer:
  name: Momentum
  momentum: 0.9
  lr:
    name: Piecewise
    learning_rate: 0.015
    decay_epochs: [30, 60, 90]
    values: [0.1, 0.01, 0.001, 0.0001]
  regularizer:
    name: 'L2'
    coeff: 0.0005

# data loader for train and eval
DataLoader:
  Train:
    dataset:
      name: ImageNetDataset
      # 根路径
      image_root: ./dataset/
      # 前面自己生产得到的训练集文本路径
      cls_label_path: ./dataset/foods/train_list.txt
      # 数据预处理
      transform_ops:
        - DecodeImage:
            to_rgb: True
            channel_first: False
        - ResizeImage:
            resize_short: 256
        - CropImage:
            size: 224
        - RandFlipImage:
            flip_code: 1
        - NormalizeImage:
            scale: 1.0/255.0
            mean: [0.485, 0.456, 0.406]
            std: [0.229, 0.224, 0.225]
            order: ''

    sampler:
      name: DistributedBatchSampler
      batch_size: 128
      drop_last: False
      shuffle: True
    loader:
      num_workers: 0
      use_shared_memory: True

  Eval:
    dataset:
      name: ImageNetDataset
      # 根路径
      image_root: ./dataset/
      # 前面自己生产得到的验证集文本路径
      cls_label_path: ./dataset/foods/val_list.txt
      # 数据预处理
      transform_ops:
        - DecodeImage:
            to_rgb: True
            channel_first: False
        - ResizeImage:
            resize_short: 256
        - CropImage:
            size: 224
        - NormalizeImage:
            scale: 1.0/255.0
            mean: [0.485, 0.456, 0.406]
            std: [0.229, 0.224, 0.225]
            order: ''
    sampler:
      name: DistributedBatchSampler
      batch_size: 128
      drop_last: False
      shuffle: True
    loader:
      num_workers: 0
      use_shared_memory: True

Infer:
  infer_imgs: ./dataset/scenes/river/66a052f90681bfca33c4e3ccd4353b105e7ee222.jpg #这里改了预测的图片路径
  batch_size: 10
  transforms:
    - DecodeImage:
        to_rgb: True
        channel_first: False
    - ResizeImage:
        resize_short: 256
    - CropImage:
        size: 224
    - NormalizeImage:
        scale: 1.0/255.0
        mean: [0.485, 0.456, 0.406]
        std: [0.229, 0.224, 0.225]
        order: ''
    - ToCHWImage:
  PostProcess:
    name: Topk
    # 输出的可能性最高的前topk个
    topk: 5
    # 标签文件 需要自己新建文件
    class_id_map_file: ./dataset/label_list.txt

Metric:
  Train:
    - TopkAcc:
        topk: [1, 5]
  Eval:
    - TopkAcc:
        topk: [1, 5]
```
#### 3.1.2 标签文件
这个是在预测时生成对照的依据，在上个文件有提到这个
```
# 标签文件 需要自己新建文件
    class_id_map_file: dataset/label_list.txt
```
按照对应的进行编写：

![](https://ai-studio-static-online.cdn.bcebos.com/0e40a0afaa824ba9b70778aa7931a3baf2a421bcb81b4b0f83632da4e4ddc0ef)

如食品分类(要对照之前的txt的类别确认无误)
```
1 beef_carpaccio
2 baby_back_ribs
3 beef_tartare
4 apple_pie
5 baklava
```
![](https://ai-studio-static-online.cdn.bcebos.com/6b3a4a244ed34517bcf4be5fcc629ab10d081eb0af7c4532a795583d19939a82)

![](https://ai-studio-static-online.cdn.bcebos.com/8ee69dbad8bf4e0f9c743645a4438cc035446d052d4b409a88c1cc25b58dcf83)

### 3.2 开始训练

```python
# 提示，运行过程中可能存在坏图的情况，但是不用担心，训练过程不受影响。
# 仅供参考，我只跑了五轮，准确率很低
!python3 tools/train.py \
    -c ./ppcls/configs/quick_start/new_user/ShuffleNetV2_x0_25.yaml
```

    [2021/08/15 14:45:06] root INFO: Already save model in ./output1/ResNet50/latest

### 3.3 预测一张

```python
# 更换为你训练的网络，需要预测的文件，上面训练所得到的的最优模型文件
# 我这里是不严谨的，直接使用训练集的图片进行验证，大家可以去百度搜一些相关的图片传上来，进行预测
!python3 tools/infer.py \
    -c ./ppcls/configs/quick_start/new_user/ShuffleNetV2_x0_25.yaml \
    -o Infer.infer_imgs=./dataset/scenes/lawn/c7a435c895f7a31ed563ba37e461d94a53a261c3.jpg \
#-o Infer.infer_imgs=./dataset/scenes/desert/9e498d798a71cb57e6c93b4d5272df10576a33e6.jpg \
    -o Global.pretrained_model=output1/ResNet50/best_model
```

    /home/aistudio/PaddleClas/ppcls/arch/backbone/model_zoo/vision_transformer.py:15: DeprecationWarning: Using or importing the ABCs from 'collections' instead of from 'collections.abc' is deprecated, and in 3.8 it will stop working
      from collections import Callable
    [2021/08/15 14:45:36] root INFO:
    ===========================================================
    ==        PaddleClas is powered by PaddlePaddle !        ==
    ===========================================================
    ==                                                       ==
    ==   For more info please go to the following website.   ==
    ==                                                       ==
    ==       https://github.com/PaddlePaddle/PaddleClas      ==
    ===========================================================

    [2021/08/15 14:45:36] root INFO: Arch :
    [2021/08/15 14:45:36] root INFO:     class_num : 6
    [2021/08/15 14:45:36] root INFO:     name : ResNet50
    [2021/08/15 14:45:36] root INFO: DataLoader :
    [2021/08/15 14:45:36] root INFO:     Eval :
    [2021/08/15 14:45:36] root INFO:         dataset :
    [2021/08/15 14:45:36] root INFO:             cls_label_path : ./dataset/scenes/val_list.txt
    [2021/08/15 14:45:36] root INFO:             image_root : ./dataset/
    [2021/08/15 14:45:36] root INFO:             name : ImageNetDataset
    [2021/08/15 14:45:36] root INFO:             transform_ops :
    [2021/08/15 14:45:36] root INFO:                 DecodeImage :
    [2021/08/15 14:45:36] root INFO:                     channel_first : False
    [2021/08/15 14:45:36] root INFO:                     to_rgb : True
    [2021/08/15 14:45:36] root INFO:                 ResizeImage :
    [2021/08/15 14:45:36] root INFO:                     resize_short : 256
    [2021/08/15 14:45:36] root INFO:                 CropImage :
    [2021/08/15 14:45:36] root INFO:                     size : 224
    [2021/08/15 14:45:36] root INFO:                 NormalizeImage :
    [2021/08/15 14:45:36] root INFO:                     mean : [0.485, 0.456, 0.406]
    [2021/08/15 14:45:36] root INFO:                     order :
    [2021/08/15 14:45:36] root INFO:                     scale : 1.0/255.0
    [2021/08/15 14:45:36] root INFO:                     std : [0.229, 0.224, 0.225]
    [2021/08/15 14:45:36] root INFO:         loader :
    [2021/08/15 14:45:36] root INFO:             num_workers : 0
    [2021/08/15 14:45:36] root INFO:             use_shared_memory : True
    [2021/08/15 14:45:36] root INFO:         sampler :
    [2021/08/15 14:45:36] root INFO:             batch_size : 128
    [2021/08/15 14:45:36] root INFO:             drop_last : False
    [2021/08/15 14:45:36] root INFO:             name : DistributedBatchSampler
    [2021/08/15 14:45:36] root INFO:             shuffle : True
    [2021/08/15 14:45:36] root INFO:     Train :
    [2021/08/15 14:45:36] root INFO:         dataset :
    [2021/08/15 14:45:36] root INFO:             cls_label_path : ./dataset/scenes/train_list.txt
    [2021/08/15 14:45:36] root INFO:             image_root : ./dataset/
    [2021/08/15 14:45:36] root INFO:             name : ImageNetDataset
    [2021/08/15 14:45:36] root INFO:             transform_ops :
    [2021/08/15 14:45:36] root INFO:                 DecodeImage :
    [2021/08/15 14:45:36] root INFO:                     channel_first : False
    [2021/08/15 14:45:36] root INFO:                     to_rgb : True
    [2021/08/15 14:45:36] root INFO:                 ResizeImage :
    [2021/08/15 14:45:36] root INFO:                     resize_short : 256
    [2021/08/15 14:45:36] root INFO:                 CropImage :
    [2021/08/15 14:45:36] root INFO:                     size : 224
    [2021/08/15 14:45:36] root INFO:                 RandFlipImage :
    [2021/08/15 14:45:36] root INFO:                     flip_code : 1
    [2021/08/15 14:45:36] root INFO:                 NormalizeImage :
    [2021/08/15 14:45:36] root INFO:                     mean : [0.485, 0.456, 0.406]
    [2021/08/15 14:45:36] root INFO:                     order :
    [2021/08/15 14:45:36] root INFO:                     scale : 1.0/255.0
    [2021/08/15 14:45:36] root INFO:                     std : [0.229, 0.224, 0.225]
    [2021/08/15 14:45:36] root INFO:         loader :
    [2021/08/15 14:45:36] root INFO:             num_workers : 0
    [2021/08/15 14:45:36] root INFO:             use_shared_memory : True
    [2021/08/15 14:45:36] root INFO:         sampler :
    [2021/08/15 14:45:36] root INFO:             batch_size : 128
    [2021/08/15 14:45:36] root INFO:             drop_last : False
    [2021/08/15 14:45:36] root INFO:             name : DistributedBatchSampler
    [2021/08/15 14:45:36] root INFO:             shuffle : True
    [2021/08/15 14:45:36] root INFO: Global :
    [2021/08/15 14:45:36] root INFO:     checkpoints : None
    [2021/08/15 14:45:36] root INFO:     device : gpu
    [2021/08/15 14:45:36] root INFO:     epochs : 30
    [2021/08/15 14:45:36] root INFO:     eval_during_train : True
    [2021/08/15 14:45:36] root INFO:     eval_interval : 1
    [2021/08/15 14:45:36] root INFO:     image_shape : [3, 224, 224]
    [2021/08/15 14:45:36] root INFO:     output_dir : ./output1/
    [2021/08/15 14:45:36] root INFO:     pretrained_model : None
    [2021/08/15 14:45:36] root INFO:     print_batch_step : 1
    [2021/08/15 14:45:36] root INFO:     save_inference_dir : ./inference
    [2021/08/15 14:45:36] root INFO:     save_interval : 1
    [2021/08/15 14:45:36] root INFO:     to_static : False
    [2021/08/15 14:45:36] root INFO:     use_visualdl : True
    [2021/08/15 14:45:36] root INFO: Infer :
    [2021/08/15 14:45:36] root INFO:     PostProcess :
    [2021/08/15 14:45:36] root INFO:         class_id_map_file : ./dataset/label_list.txt
    [2021/08/15 14:45:36] root INFO:         name : Topk
    [2021/08/15 14:45:36] root INFO:         topk : 5
    [2021/08/15 14:45:36] root INFO:     batch_size : 10
    [2021/08/15 14:45:36] root INFO:     infer_imgs : ./dataset/scenes/lawn/c7a435c895f7a31ed563ba37e461d94a53a261c3.jpg
    [2021/08/15 14:45:36] root INFO:     transforms :
    [2021/08/15 14:45:36] root INFO:         DecodeImage :
    [2021/08/15 14:45:36] root INFO:             channel_first : False
    [2021/08/15 14:45:36] root INFO:             to_rgb : True
    [2021/08/15 14:45:36] root INFO:         ResizeImage :
    [2021/08/15 14:45:36] root INFO:             resize_short : 256
    [2021/08/15 14:45:36] root INFO:         CropImage :
    [2021/08/15 14:45:36] root INFO:             size : 224
    [2021/08/15 14:45:36] root INFO:         NormalizeImage :
    [2021/08/15 14:45:36] root INFO:             mean : [0.485, 0.456, 0.406]
    [2021/08/15 14:45:36] root INFO:             order :
    [2021/08/15 14:45:36] root INFO:             scale : 1.0/255.0
    [2021/08/15 14:45:36] root INFO:             std : [0.229, 0.224, 0.225]
    [2021/08/15 14:45:36] root INFO:         ToCHWImage : None
    [2021/08/15 14:45:36] root INFO: Loss :
    [2021/08/15 14:45:36] root INFO:     Eval :
    [2021/08/15 14:45:36] root INFO:         CELoss :
    [2021/08/15 14:45:36] root INFO:             weight : 1.0
    [2021/08/15 14:45:36] root INFO:     Train :
    [2021/08/15 14:45:36] root INFO:         CELoss :
    [2021/08/15 14:45:36] root INFO:             weight : 1.0
    [2021/08/15 14:45:36] root INFO: Metric :
    [2021/08/15 14:45:36] root INFO:     Eval :
    [2021/08/15 14:45:36] root INFO:         TopkAcc :
    [2021/08/15 14:45:36] root INFO:             topk : [1, 5]
    [2021/08/15 14:45:36] root INFO:     Train :
    [2021/08/15 14:45:36] root INFO:         TopkAcc :
    [2021/08/15 14:45:36] root INFO:             topk : [1, 5]
    [2021/08/15 14:45:36] root INFO: Optimizer :
    [2021/08/15 14:45:36] root INFO:     lr :
    [2021/08/15 14:45:36] root INFO:         decay_epochs : [30, 60, 90]
    [2021/08/15 14:45:36] root INFO:         learning_rate : 0.015
    [2021/08/15 14:45:36] root INFO:         name : Piecewise
    [2021/08/15 14:45:36] root INFO:         values : [0.1, 0.01, 0.001, 0.0001]
    [2021/08/15 14:45:36] root INFO:     momentum : 0.9
    [2021/08/15 14:45:36] root INFO:     name : Momentum
    [2021/08/15 14:45:36] root INFO:     regularizer :
    [2021/08/15 14:45:36] root INFO:         coeff : 0.0005
    [2021/08/15 14:45:36] root INFO:         name : L2
    W0815 14:45:36.183862 17099 device_context.cc:404] Please NOTE: device: 0, GPU Compute Capability: 7.0, Driver API Version: 10.1, Runtime API Version: 10.1
    W0815 14:45:36.189221 17099 device_context.cc:422] device: 0, cuDNN Version: 7.6.
    [2021/08/15 14:45:41] root INFO: train with paddle 2.1.2 and device CUDAPlace(0)
    /opt/conda/envs/python35-paddle120-env/lib/python3.7/site-packages/paddle/tensor/creation.py:125: DeprecationWarning: `np.object` is a deprecated alias for the builtin `object`. To silence this warning, use `object` by itself. Doing this will not modify any behavior and is safe.
    Deprecated in NumPy 1.20; for more details and guidance: https://numpy.org/devdocs/release/1.20.0-notes.html#deprecations
      if data.dtype == np.object:
    [{'class_ids': [5, 2, 4, 3, 1], 'scores': [1.0, 0.0, 0.0, 0.0, 0.0], 'file_name': './dataset/scenes/lawn/c7a435c895f7a31ed563ba37e461d94a53a261c3.jpg', 'label_names': ['desert', 'lawn', 'church', 'river', 'ice']}]

运行完成，最后几行会得到结果如下形式：
```
[{'class_ids': [5, 1, 3, 4, 2],
'scores': [0.48433, 0.26765, 0.13903, 0.05609, 0.05162],
'file_name': 'dataset/foods/baby_back_ribs/319516.jpg',
'label_names': ['baklava', 'beef_carpaccio', 'beef_tartare', 'apple_pie', 'baby_back_ribs']}]
```
可以发现，预测结果不对，准确率很低，但整体的项目流程你已经掌握了！
训练轮数还有很大提升空间，自行变动参数直到预测正确为止~
这里最终我训练出来的模型好像有点问题，scores有一个永远是1.0，其他是0.0，这种现象应该是有问题的，而且分类有些错误，后续还得再进行更改！
恭喜你学会了paddleclas图像分类！

## 最后（！很重要！）
1. 生成版本
![](https://ai-studio-static-online.cdn.bcebos.com/adda3ac3b96c403daa68f9f5ec9fd807930d1404f5834317a4f3d7e662243ebc)
2. 有能力者导出markdown发表到github上
![](https://ai-studio-static-online.cdn.bcebos.com/a1f8fdad7659463381d8857e9fd63e84648591fc1a264f1a8241029fb825fe78)
3. 公开项目
![](https://ai-studio-static-online.cdn.bcebos.com/bc8af42d84914b7a8e882e3b88551aa7a5943cf7351747b2b6df3310be53eb05)
4. 祝大家顺利结业~


"""Tasks for beam measurement processing"""
import json
import yaml
from os import path, listdir

import numpy as np
from scipy import constants

from caput import config, tod, mpiarray, mpiutil
from caput.pipeline import PipelineRuntimeError
from caput.time import STELLAR_S

from ch_util import ephemeris as ephem
from ch_util import tools, layout, holography
from chimedb import data_index as di
from chimedb.core import connect as connect_database

from draco.core import task, io
from draco.util import regrid
from draco.analysis.transform import Regridder
from draco.core.containers import ContainerBase, SiderealStream, TimeStream, TrackBeam
from draco.util.tools import invert_no_zero

from ..core.containers import TransitFitParams
from .calibration import TransitFit, GainFromTransitFit

SIDEREAL_DAY_SEC = STELLAR_S * 24 * 3600
SPEED_LIGHT = float(constants.c) / 1e6  # 10^6 m / s
CHIME_CYL_W = 22.0  # m

class TransitGrouper(task.SingleTask):
    """Group transits from a sequence of TimeStream objects.

    Attributes
    ----------
    ha_span: float
        Span in degrees surrounding transit.
    min_span: float
        Minimum span (deg) of a transit to accept.
    source: str
        Name of the transiting source. (Must match what is used in `ch_util.ephemeris`.)
    db_source: str
        Name of the transiting source as listed in holography database.
        This is a hack until a better solution is implemented.
    """

    ha_span = config.Property(proptype=float, default=180.0)
    min_span = config.Property(proptype=float, default=0.0)
    source = config.Property(proptype=str)
    db_source = config.Property(proptype=str)

    def setup(self, observer=None):
        """Set the local observers position if not using CHIME.

        Parameters
        ----------
        observer : caput.time.Observer, optional
            Details of the observer, if not set default to CHIME.

        Raises
        ------
        `caput.config.CaputConfigError`
            If config value ``source`` doesn't match any in `ch_util.ephemeris`.
        """
        self.observer = ephem.chime if observer is None else observer
        try:
            self.src = ephem.source_dictionary[self.source]
        except KeyError:
            msg = (
                "Could not find source {} in catalogue. "
                "Must use same spelling as in `ch_util.ephemeris`.".format(self.source)
            )
            self.log.error(msg)
            raise config.CaputConfigError(msg)
        self.cur_transit = None
        self.tstreams = []
        self.last_time = 0

        # Get list of holography observations
        # Only allowed to query database from rank0
        db_runs = None
        if mpiutil.rank0:
            connect_database()
            db_runs = list(get_holography_obs(self.db_source))
            db_runs = [(int(r.id), (r.start_time, r.finish_time)) for r in db_runs]
        self.db_runs = mpiutil.bcast(db_runs, root=0)
        mpiutil.barrier()

    def process(self, tstream):
        """Take in a timestream and accumulate, group into whole transits.

        Parameters
        ----------
        tstream : TimeStream

        Returns
        -------
        ts : TimeStream
            Timestream with containing transits of specified length.
        """

        # Redistribute if needed
        tstream.redistribute("freq")

        # placeholder for finalized transit when it is ready
        final_ts = None

        # check if we jumped to another acquisition
        if (tstream.time[0] - self.last_time) > 5 * (tstream.time[1] - tstream.time[0]):
            if self.cur_transit is None:
                # will be true when we start a new transit
                pass
            else:
                # start on a new transit and return current one
                final_ts = self._finalize_transit()

        # if this is the start of a new grouping, setup transit bounds
        if self.cur_transit is None:
            tr_time = self.observer.transit_times(self.src, tstream.time[0])
            if len(tr_time) != 1:
                raise ValueError(
                    "Didn't find exactly one transit time. Found {:d}.".format(
                        len(tr_time)
                    )
                )
            self.cur_transit = tr_time[0]
            self._transit_bounds()

        # check if we've accumulated enough past the transit
        if tstream.time[-1] > self.end_t:
            self.append(tstream)
            final_ts = self._finalize_transit()
        elif tstream.time[-1] < self.start_t:
            pass
        else:
            self.append(tstream)

        self.last_time = tstream.time[-1]

        return final_ts

    def process_finish(self):
        """Return the current transit before finishing.

        Returns
        -------
        ts: TimeStream
            Last (possibly incomplete) transit.
        """
        return self._finalize_transit()

    def append(self, ts):
        """Append a timestream to the buffer list.

        This will strip eigenvector datasets if they are present.
        """
        for dname in ["evec", "eval", "erms"]:
            if dname in ts.datasets.keys():
                self.log.debug("Stripping dataset {}".format(dname))
                del ts[dname]
        self.tstreams.append(ts)

    def _finalize_transit(self):
        """Concatenate grouped time streams for the currrent transit."""

        # Find where transit starts and ends
        if len(self.tstreams) == 0 or self.cur_transit is None:
            self.log.info("Did not find any transits.")
            return None
        self.log.debug(
            "Finalising transit for {}...".format(
                ephem.unix_to_datetime(self.cur_transit)
            )
        )
        all_t = np.concatenate([ts.time for ts in self.tstreams])
        start_ind = int(np.argmin(np.abs(all_t - self.start_t)))
        stop_ind = int(np.argmin(np.abs(all_t - self.end_t)))

        # Save list of filenames
        filenames = [ts.attrs["filename"] for ts in self.tstreams]

        dt = self.tstreams[0].time[1] - self.tstreams[0].time[0]
        if dt <= 0:
            self.log.warning(
                "Time steps are not positive definite: dt={:.3f}".format(dt)
                + " Skipping."
            )
            ts = None
        if stop_ind - start_ind > int(self.min_span / 360.0 * SIDEREAL_DAY_SEC / dt):
            if len(self.tstreams) > 1:
                # Concatenate timestreams
                ts = tod.concatenate(self.tstreams, start=start_ind, stop=stop_ind)
            else:
                ts = self.tstreams[0]
            _, dec = ephem.object_coords(
                self.src, all_t[0], deg=True, obs=self.observer
            )
            ts.attrs["dec"] = dec
            ts.attrs["source_name"] = self.source
            ts.attrs["transit_time"] = self.cur_transit
            ts.attrs["observation_id"] = self.obs_id
            ts.attrs["tag"] = "{}_{:0>4d}_{}".format(
                self.source,
                self.obs_id,
                ephem.unix_to_datetime(self.cur_transit).strftime("%Y%m%dT%H%M%S"),
            )
            ts.attrs["archivefiles"] = filenames
        else:
            self.log.info("Transit too short. Skipping.")
            ts = None

        self.tstreams = []
        self.cur_transit = None

        return ts

    def _transit_bounds(self):
        """Find the start and end times of this transit.

        Compares the desired HA span to the start and end times of the observation
        recorded in the database. Also gets the observation ID."""

        # subtract half a day from start time to ensure we don't get following day
        self.start_t = self.cur_transit - self.ha_span / 360.0 / 2.0 * SIDEREAL_DAY_SEC
        self.end_t = self.cur_transit + self.ha_span / 360.0 / 2.0 * SIDEREAL_DAY_SEC

        # get bounds of observation from database
        this_run = [
            r
            for r in self.db_runs
            if r[1][0] < self.cur_transit and r[1][1] > self.cur_transit
        ]
        if len(this_run) == 0:
            self.log.warning(
                "Could not find source transit in holography database for {}.".format(
                    ephem.unix_to_datetime(self.cur_transit)
                )
            )
            # skip this file
            self.cur_transit = None
        else:
            self.start_t = max(self.start_t, this_run[0][1][0])
            self.end_t = min(self.end_t, this_run[0][1][1])
            self.obs_id = this_run[0][0]

class TransitRegridder(Regridder):
    """Interpolate TimeStream transits onto a regular grid in hour angle.

    Attributes
    ----------
    samples : int
        Number of samples to interpolate onto.
    ha_span: float
        Span in degrees surrounding transit.
    lanczos_width : int
        Width of the Lanczos interpolation kernel.
    snr_cov: float
        Ratio of signal covariance to noise covariance (used for Wiener filter).
    source: str
        Name of the transiting source. (Must match what is used in `ch_util.ephemeris`.)
    """

    samples = config.Property(proptype=int, default=1024)
    lanczos_width = config.Property(proptype=int, default=5)
    snr_cov = config.Property(proptype=float, default=1e-8)
    ha_span = config.Property(proptype=float, default=180.0)
    source = config.Property(proptype=str)

    def setup(self, observer=None):
        """Set the local observers position if not using CHIME.

        Parameters
        ----------
        observer : caput.time.Observer, optional
            Details of the observer, if not set default to CHIME.

        Raises
        ------
        `caput.config.CaputConfigError`
            If config value ``source`` doesn't match any in `ch_util.ephemeris`.
        """
        self.observer = ephem.chime if observer is None else observer

        # Setup bounds for interpolation grid
        self.start = -self.ha_span / 2
        self.end = self.ha_span / 2

        try:
            self.src = ephem.source_dictionary[self.source]
        except KeyError:
            msg = (
                "Could not find source {} in catalogue. "
                "Must use same spelling as in `ch_util.ephemeris`.".format(self.source)
            )
            self.log.error(msg)
            raise config.CaputConfigError(msg)

    def process(self, data):
        """Regrid visibility data onto a regular grid in hour angle.

        Parameters
        ----------
        data : TimeStream
            Time-ordered data.

        Returns
        -------
        new_data : SiderealStream
            The regridded data centered on the source RA.
        """

        # Redistribute if needed
        data.redistribute("freq")

        # View of data
        weight = data.weight[:].view(np.ndarray)
        vis_data = data.vis[:].view(np.ndarray)

        # Get apparent source RA, including precession effects
        ra, _ = ephem.object_coords(self.src, data.time[0], deg=True, obs=self.observer)
        # Get catalogue RA for reference
        ra_icrs, _ = ephem.object_coords(self.src, deg=True, obs=self.observer)

        # Convert input times to hour angle
        lha = unwrap_lha(self.observer.unix_to_lsa(data.time), ra)

        # perform regridding
        success = 1
        try:
            new_grid, new_vis, ni = self._regrid(vis_data, weight, lha)
        except np.linalg.LinAlgError as e:
            self.log.error(str(e))
            success = 0
        except ValueError as e:
            self.log.error(str(e))
            success = 0
        # Check other ranks have completed
        success = mpiutil.allreduce(success)
        if success != mpiutil.size:
            self.log.warning("Regridding failed. Skipping transit.")
            return None

        # mask out regions beyond bounds of this transit
        grid_mask = np.ones_like(new_grid)
        grid_mask[new_grid < lha.min()] = 0.0
        grid_mask[new_grid > lha.max()] = 0.0
        new_vis *= grid_mask
        ni *= grid_mask

        # Wrap to produce MPIArray
        if data.distributed:
            new_vis = mpiarray.MPIArray.wrap(new_vis, axis=data.vis.distributed_axis)
            ni = mpiarray.MPIArray.wrap(ni, axis=data.vis.distributed_axis)

        # Create new container for output
        ra_grid = (new_grid + ra) % 360.0
        new_data = SiderealStream(
            axes_from=data, attrs_from=data, ra=ra_grid, comm=data.comm
        )
        new_data.redistribute("freq")
        new_data.vis[:] = new_vis
        new_data.weight[:] = ni
        new_data.attrs["cirs_ra"] = ra
        new_data.attrs["icrs_ra"] = ra_icrs

        return new_data

class EdgeFlagger(task.SingleTask):
    """Flag the edges of the transit.

    Parameters
    ----------
    num_begin: int
        Number of samples to flag at the start of the transit.
    num_end: int
        Number of samples to flag at the end of the transit.
    """

    num_begin = config.Property(proptype=int, default=15)
    num_end = config.Property(proptype=int, default=15)

    def process(self, track):
        """Extend the region that has weight set to zero.

        Parameters
        ----------
        track: draco.core.containers.TrackBeam
            Holography track to flag.

        Returns
        -------
        track: draco.core.containers.TrackBeam
            Holography track with weight set to zero for
            `num_begin` samples at the start of the transit
            and `num_end` samples at the end.
        """
        if (self.num_begin == 0) and (self.num_end == 0):
            return track

        track.redistribute("freq")

        weight = track["weight"][:].view(np.ndarray)

        for ind in np.ndindex(*weight.shape[:-1]):

            flag = np.flatnonzero(weight[ind] > 0.0)

            if flag.size > 0:

                imin, imax = np.percentile(flag, [0, 100]).astype(np.int)
                imax = imax + 1

                if self.num_begin:
                    weight[ind][imin : imin + self.num_begin] = 0.0

                if self.num_end:
                    weight[ind][imax - self.num_end : imax] = 0.0

        return track

class TransitResampler(task.SingleTask):
    """Resample the beam at specific RAs.

    Attributes
    ----------
    lanczos_width : int
        Width of the Lanczos kernel in number of samples.
    """

    lanczos_width = config.Property(proptype=int, default=5)

    def setup(self, observer=None):
        """Set the local observers position if not using CHIME.

        Parameters
        ----------
        observer : caput.time.Observer, optional
            Details of the observer, if not set default to CHIME.
        """
        self.observer = ephem.chime if observer is None else observer

    def process(self, beam, data):
        """Resample the measured beam at arbitrary RA by convolving with a Lanczos kernel.

        Parameters
        ----------
        beam : draco.core.containers.TrackBeam
            The beam that will be resampled.
        data : draco.core.containers.VisContainer
            Must contain a `time` or `ra` axis that defines
            the RAs to sample.

        Returns
        -------
        new_beam : draco.core.containers.TrackBeam
             The input `beam` re-sampled at the RAs contained in `data`.
        """
        # Distribute over frequencies
        data.redistribute("freq")
        beam.redistribute("freq")

        # Dereference datasets
        bv = beam.beam[:].view(np.ndarray)
        bw = beam.weight[:].view(np.ndarray)

        flag = np.any(bw > 0.0, axis=(0, 1, 2))

        # Compute the hour angle of the source
        if "ra" in data.index_map:
            ra = data.index_map["ra"][:]
        elif "time" in data.index_map:
            ra = self.observer.unix_to_lsa(data.time)
        else:
            raise RuntimeError("Unable to extract RA from input container.")

        hour_angle = beam.pix["phi"][:]

        new_hour_angle = (((ra - beam.attrs["cirs_ra"]) + 180.0) % 360.0) - 180.0

        isort = np.argsort(new_hour_angle)
        new_hour_angle_sorted = new_hour_angle[isort]

        within_range = np.flatnonzero(
            (new_hour_angle_sorted >= np.min(hour_angle[flag]))
            & (new_hour_angle_sorted <= np.max(hour_angle[flag]))
        )

        if within_range.size < (2 * self.lanczos_width):
            raise RuntimeError("Not enough overlapping samples.")

        slc = slice(
            np.min(within_range) + self.lanczos_width,
            np.max(within_range) - self.lanczos_width + 1,
        )

        # Perform regridding
        new_bv, new_bw = self._resample(hour_angle, bv, bw, new_hour_angle_sorted[slc])

        # Create output container
        new_beam = TrackBeam(
            phi=new_hour_angle,
            theta=np.repeat(np.median(beam.pix["theta"][:]), new_hour_angle.size),
            axes_from=beam,
            attrs_from=beam,
            distributed=beam.distributed,
            comm=beam.comm,
        )

        new_beam.redistribute("freq")

        # Save to container
        new_beam.beam[:] = 0.0
        new_beam.weight[:] = 0.0

        new_beam.beam[:, :, :, isort[slc]] = new_bv
        new_beam.weight[:, :, :, isort[slc]] = new_bw

        return new_beam

    def _resample(self, xgrid, ygrid, wgrid, x):

        lza = regrid.lanczos_forward_matrix(xgrid, x, a=self.lanczos_width).T

        y = np.matmul(ygrid, lza)
        w = invert_no_zero(np.matmul(invert_no_zero(wgrid), lza ** 2))

        return y, w

class MakeHolographyBeam(task.SingleTask):
    """Repackage a holography transit into a beam container.

    The visibilities will be grouped according to their respective 26 m
    input (along the `pol` axis, labelled by the 26 m polarisation of that input).

    The form of the beam dataset is A_{i} = < E_{i} * E_{26m}^{*} >, i.e. the 26m
    input is the conjugate part of the visbility.
    """

    def process(self, data, inputmap):
        """Package a holography transit into a beam container.

        Parameters
        ----------
        data : SiderealStream
            Transit observation as generated by `TransitRegridder`.
        inputmap : list of `CorrInput`
            A list describing the inputs as they are in the file, output from
            `ch_util.tools.get_correlator_inputs()`

        Returns
        -------
        track : TrackBeam
            The transit in a beam container.
        """

        # redistribute if needed
        data.redistribute("freq")

        prod = data.index_map["prod"]
        inputs = data.index_map["input"]

        # Figure out which inputs are the 26m
        input_26m = prod["input_a"][np.where(prod["input_a"] == prod["input_b"])[0]]
        if len(input_26m) != 2:
            msg = "Did not find exactly two 26m inputs in the data."
            self.log.error(msg)
            raise PipelineRuntimeError(msg)

        # Separate products by 26 m inputs
        prod_groups = []
        for i in input_26m:
            prod_groups.append(
                np.where(np.logical_or(prod["input_a"] == i, prod["input_b"] == i))[0]
            )

        # Check we have the expected number of products
        if (
            prod_groups[0].shape[0] != inputs.shape[0]
            or prod_groups[1].shape[0] != inputs.shape[0]
        ):
            msg = (
                "Products do not separate into two groups with the length of the input map. "
                "({:d}, {:d}) != {:d}"
            ).format(prod_groups[0].shape[0], prod_groups[1].shape[0], inputs.shape[0])
            self.log.error(msg)
            raise PipelineRuntimeError(msg)

        # Sort based on the id in the layout database
        corr_id = np.array([inp.id for inp in inputmap])
        isort = np.argsort(corr_id)

        # Create new input axis using id and serial number in database
        inputs_sorted = np.array(
            [(inputmap[ii].id, inputmap[ii].input_sn) for ii in isort],
            dtype=inputs.dtype,
        )

        # Sort the products based on the input id in database and
        # determine which products should be conjugated.
        conj = []
        prod_groups_sorted = []
        for i, pg in enumerate(prod_groups):
            group_prod = prod[pg]
            group_conj = group_prod["input_a"] == input_26m[i]
            group_inputs = np.where(
                group_conj, group_prod["input_b"], group_prod["input_a"]
            )
            group_sort = np.argsort(corr_id[group_inputs])

            prod_groups_sorted.append(pg[group_sort])
            conj.append(group_conj[group_sort])

        # Regroup by co/cross-pol
        copol, xpol = [], []
        prod_groups_cox = [pg.copy() for pg in prod_groups_sorted]
        conj_cox = [pg.copy() for pg in conj]
        input_pol = np.array(
            [
                ipt.pol
                if (tools.is_array(ipt) or tools.is_holographic(ipt))
                else inputmap[input_26m[0]].pol
                for ipt in inputmap
            ]
        )
        for i, pg in enumerate(prod_groups_sorted):
            group_prod = prod[pg]
            # Determine co/cross in each prod group
            cp = (
                input_pol[
                    np.where(conj[i], group_prod["input_b"], group_prod["input_a"])
                ]
                == inputmap[input_26m[i]].pol
            )
            xp = np.logical_not(cp)
            copol.append(cp)
            xpol.append(xp)
            # Move products to co/cross-based groups
            prod_groups_cox[0][cp] = pg[cp]
            prod_groups_cox[1][xp] = pg[xp]
            conj_cox[0][cp] = conj[i][cp]
            conj_cox[1][xp] = conj[i][xp]
        # Check for compeleteness
        consistent = np.all(copol[0] + copol[1] == np.ones(copol[0].shape)) and np.all(
            xpol[0] + xpol[1] == np.ones(xpol[0].shape)
        )
        if not consistent:
            msg = (
                "Products do not separate exclusively into co- and cross-polar groups."
            )
            self.log.error(msg)
            raise PipelineRuntimeError(msg)

        # Make new index map
        ra = data.attrs["cirs_ra"]
        phi = unwrap_lha(data.ra[:], ra)
        if "dec" not in data.attrs.keys():
            msg = (
                "Input stream must have a 'dec' attribute specifying "
                "declination of holography source."
            )
            self.log.error(msg)
            raise PipelineRuntimeError(msg)
        theta = np.ones_like(phi) * data.attrs["dec"]
        pol = np.array(["co", "cross"], dtype="S5")

        # Create new container and fill
        track = TrackBeam(
            theta=theta,
            phi=phi,
            track_type="drift",
            coords="celestial",
            input=inputs_sorted,
            pol=pol,
            freq=data.freq[:],
            attrs_from=data,
            distributed=data.distributed,
        )
        for ip in range(len(pol)):
            track.beam[:, ip, :, :] = data.vis[:, prod_groups_cox[ip], :]
            track.weight[:, ip, :, :] = data.weight[:, prod_groups_cox[ip], :]
            if np.any(conj_cox[ip]):
                track.beam[:, ip, conj_cox[ip], :] = track.beam[
                    :, ip, conj_cox[ip], :
                ].conj()

        # Store 26 m inputs
        track.attrs["26m_inputs"] = [list(isort).index(ii) for ii in input_26m]

        return track

class ConstructStackedBeam(task.SingleTask):
    """Construct the effective beam for stacked baselines.

    Parameters
    ----------
    weight : string ('uniform' or 'inverse_variance')
        How to weight the baselines when stacking:
            'uniform' - each baseline given equal weight
            'inverse_variance' - each baseline weighted by the weight attribute
    """

    weight = config.enum(["uniform", "inverse_variance"], default="uniform")

    def setup(self, tel):
        """Set the Telescope instance to use.

        Parameters
        ----------
        tel : TransitTelescope
        """
        self.telescope = io.get_telescope(tel)

    def process(self, beam, data):
        """Stack

        Parameters
        ----------
        beam : TrackBeam
            The beam that will be stacked.
        data : VisContainer
            Must contain `prod` index map and `stack` reverse map
            that will be used to stack the beam.

        Returns
        -------
        stacked_beam: VisContainer
            The input `beam` stacked in the same manner as

        """
        # Distribute over frequencies
        data.redistribute("freq")
        beam.redistribute("freq")

        # Grab the stack specifications from the input sidereal stream
        prod = data.index_map["prod"]
        reverse_stack = data.reverse_map["stack"][:]

        input_flags = data.input_flags[:]
        if not np.any(input_flags):
            input_flags = np.ones_like(input_flags)

        # Create output container
        if isinstance(data, SiderealStream):
            OutputContainer = SiderealStream
            output_kwargs = {"ra": data.ra[:]}
        else:
            OutputContainer = TimeStream
            output_kwargs = {"time": data.time[:]}

        stacked_beam = OutputContainer(
            axes_from=data,
            attrs_from=beam,
            distributed=True,
            comm=data.comm,
            **output_kwargs
        )

        stacked_beam.vis[:] = 0.0
        stacked_beam.weight[:] = 0.0

        stacked_beam.attrs["tag"] = "_".join([beam.attrs["tag"], data.attrs["tag"]])

        # Dereference datasets
        bv = beam.beam[:].view(np.ndarray)
        bw = beam.weight[:].view(np.ndarray)

        ov = stacked_beam.vis[:]
        ow = stacked_beam.weight[:]

        pol_filter = {
            "X": "X",
            "Y": "Y",
            "E": "X",
            "S": "Y",
            "co": "co",
            "cross": "cross",
        }
        pol = [pol_filter.get(pp, None) for pp in self.telescope.polarisation]
        beam_pol = [pol_filter.get(pp, None) for pp in beam.index_map["pol"][:]]

        # Compute the fractional variance of the beam measurement
        frac_var = invert_no_zero(bw * np.abs(bv) ** 2)

        # Create counter to increment during the stacking.
        # This will be used to normalize at the end.
        counter = np.zeros_like(ow)

        # Construct stack
        for pp, (ss, conj) in enumerate(reverse_stack):

            aa, bb = prod[pp]
            if conj:
                aa, bb = bb, aa

            try:
                aa_pol, bb_pol = self._resolve_pol(pol[aa], pol[bb], beam_pol)
            except ValueError:
                continue

            cross = bv[:, aa_pol, aa, :] * bv[:, bb_pol, bb, :].conj()

            weight = (
                input_flags[np.newaxis, aa, :]
                * input_flags[np.newaxis, bb, :]
                * invert_no_zero(
                    np.abs(cross) ** 2
                    * (frac_var[:, aa_pol, aa, :] + frac_var[:, bb_pol, bb, :])
                )
            )

            if self.weight == "inverse_variance":
                wss = weight
            else:
                wss = (weight > 0.0).astype(np.float32)

            # Accumulate variances in quadrature.  Save in the weight dataset.
            ov[:, ss, :] += wss * cross
            ow[:, ss, :] += wss ** 2 * invert_no_zero(weight)

            # Increment counter
            counter[:, ss, :] += wss

        # Divide through by counter to get properly weighted visibility average
        ov[:] *= invert_no_zero(counter)
        ow[:] = counter ** 2 * invert_no_zero(ow[:])

        return stacked_beam

    @staticmethod
    def _resolve_pol(pol1, pol2, pol_axis):

        if "co" in pol_axis:

            if pol1 == pol2:
                ipol = pol_axis.index("co")
            else:
                ipol = pol_axis.index("cross")

            return ipol, ipol

        else:

            if pol1 == pol2:
                ipol1 = pol_axis.index(pol1)
                ipol2 = pol_axis.index(pol2)
            else:
                ipol1 = pol_axis.index(pol2)
                ipol2 = pol_axis.index(pol1)

            return ipol1, ipol2

class HolographyTransitFit(TransitFit):
    """Fit a model to the transit.

    Attributes
    ----------
    model : str
        Name of the model to fit.  One of 'gauss_amp_poly_phase' or
        'poly_log_amp_poly_phase'.
    nsigma : float
        Number of standard deviations away from transit to fit.
    absolute_sigma : bool
        Set to True if the errors provided are absolute.  Set to False if
        the errors provided are relative, in which case the parameter covariance
        will be scaled by the chi-squared per degree-of-freedom.
    poly_type : str
        Type of polynomial.  Either 'standard', 'hermite', or 'chebychev'.
    poly_deg_amp : int
        Degree of the polynomial to fit to amplitude.
        Relevant if `poly = True`.
    poly_deg_phi : int
        Degree of the polynomial to fit to phase.
        Relevant if `poly = True`.
    niter : int
        Number of times to update the errors using model amplitude.
        Relevant if `poly = True`.
    moving_window : int
        Number of standard deviations away from peak to fit.
        The peak location is updated with each iteration.
        Must be less than `nsigma`.  Relevant if `poly = True`.
    """

    # Disable NaN checks by default since they can be valid outputs of the fit
    nan_check = config.Property(default=False, proptype=bool)
    nan_skip = config.Property(default=False, proptype=bool)
    nan_dump = config.Property(default=False, proptype=bool)

    def process(self, transit):
        """Perform the fit.

        Parameters
        ----------
        transit: TrackBeam
            Transit to be fit.

        Returns
        -------
        fit: TransitFitParams
            Fit parameters.
        """
        transit.redistribute("freq")

        # Set initial estimate of beam sigma
        local_slice = slice(
            transit.beam.local_offset[0],
            transit.beam.local_offset[0] + transit.beam.local_shape[0],
        )
        ninput = transit.beam.local_shape[2]
        freq = transit.freq[local_slice]
        sigma = (0.7 * SPEED_LIGHT / (CHIME_CYL_W * freq)) * (360.0 / np.pi)
        sigma = sigma[:, np.newaxis] * np.ones((1, ninput), dtype=sigma.dtype)

        # Find index into pol axis that yields copolar products
        pol_axis = list(transit.index_map["pol"])
        if "co" in pol_axis:
            copolar_slice = (slice(None), pol_axis.index("co"))
        else:
            this_pol = np.array(
                [
                    pol_axis.index("S")
                    if not ((ii // 256) % 2)
                    else pol_axis.index("E")
                    for ii in range(ninput)
                ]
            )
            copolar_slice = (slice(None), this_pol, np.arange(ninput))

        # Dereference datasets
        ha = transit.pix["phi"][:]

        vis = transit.beam[:].view(np.ndarray)
        vis = vis[copolar_slice]

        err = transit.weight[:].view(np.ndarray)
        err = np.sqrt(invert_no_zero(err[copolar_slice]))

        # Flag data that is outside the fit window set by nsigma config parameter
        if self.nsigma is not None:
            err *= (
                np.abs(ha[np.newaxis, np.newaxis, :])
                <= (self.nsigma * sigma[:, :, np.newaxis])
            ).astype(err.dtype)

        # Instantiate the model fitter
        model = self.ModelClass(**self.model_kwargs)

        # Fit the model
        model.fit(ha, vis, err, width=sigma, **self.fit_kwargs)

        # Pack into container
        fit = TransitFitParams(
            param=model.parameter_names,
            component=model.component,
            axes_from=transit,
            attrs_from=transit,
            distributed=transit.distributed,
            comm=transit.comm,
        )

        fit.add_dataset("chisq")
        fit.add_dataset("ndof")

        fit.redistribute("freq")

        # Transfer fit information to container attributes
        fit.attrs["model_kwargs"] = json.dumps(model.model_kwargs)
        fit.attrs["model_class"] = ".".join(
            [getattr(self.ModelClass, key) for key in ["__module__", "__name__"]]
        )

        # Save datasets
        fit.parameter[:] = model.param[:]
        fit.parameter_cov[:] = model.param_cov[:]
        fit.chisq[:] = model.chisq[:]
        fit.ndof[:] = model.ndof[:]

        return fit

class ApplyHolographyGains(task.SingleTask):
    """Apply gains to a holography transit

    Attributes
    ----------
    overwrite: bool (default: False)
        If True, overwrite the input TrackBeam.
    """

    overwrite = config.Property(proptype=bool, default=False)

    def process(self, track_in, gain):
        """Apply gain

        Parameters
        ----------
        track: draco.core.containers.TrackBeam
            Holography track to apply gains to. Will apply gains to
            track['beam'], expecting axes to be freq, pol, input, ha
        gain: np.array
            Gain to apply. Expected axes are freq, pol, and input

        Returns
        -------
        track: draco.core.containers.TrackBeam
            Holography track with gains applied.
        """

        if self.overwrite:
            track = track_in
        else:
            track = TrackBeam(
                axes_from=track_in,
                attrs_from=track_in,
                distributed=track_in.distributed,
                comm=track_in.comm,
            )
            track["beam"] = track_in["beam"][:]
            track["weight"] = track_in["weight"][:]

        track["beam"][:] *= gain.gain[:][:, np.newaxis, :, np.newaxis]
        track["weight"][:] *= invert_no_zero(np.abs(gain.gain[:]) ** 2)[
            :, np.newaxis, :, np.newaxis
        ]

        track["beam"][:] = np.where(np.isfinite(track["beam"][:]), track["beam"][:], 0)
        track["weight"][:] = np.where(
            np.isfinite(track["weight"][:]), track["weight"][:], 0
        )

        return track

class TransitStacker(task.SingleTask):
    """Stack a number of transits together.

    The weights will be inverted and stacked as variances. The variance
    between transits is evaluated and recorded in the output container.

    All transits should be on a common grid in hour angle.

    Attributes
    ----------
    weight: str (default: uniform)
        The weighting to use in the stack. One of `uniform` or `inverse_variance`.
    """

    weight = config.enum(["uniform", "inverse_variance"], default="uniform")

    def setup(self):
        """Initialise internal variables."""

        self.stack = None
        self.variance = None
        self.pseudo_variance = None
        self.norm = None

    def process(self, transit):
        """Add a transit to the stack.

        Parameters
        ----------
        transit: draco.core.containers.TrackBeam
            A holography transit.
        """

        self.log.info("Weight is %s" % self.weight)

        if self.stack is None:
            self.log.info("Initializing transit stack.")
            self.stack = TrackBeam(
                axes_from=transit, distributed=transit.distributed, comm=transit.comm
            )

            self.stack.add_dataset("sample_variance")
            self.stack.add_dataset("nsample")
            self.stack.redistribute("freq")

            self.log.info("Adding %s to stack." % transit.attrs["tag"])

            # Copy over relevant attributes
            self.stack.attrs["filename"] = [transit.attrs["tag"]]
            self.stack.attrs["observation_id"] = [transit.attrs["observation_id"]]
            self.stack.attrs["transit_time"] = [transit.attrs["transit_time"]]
            self.stack.attrs["archivefiles"] = list(transit.attrs["archivefiles"])

            self.stack.attrs["dec"] = transit.attrs["dec"]
            self.stack.attrs["source_name"] = transit.attrs["source_name"]
            self.stack.attrs["icrs_ra"] = transit.attrs["icrs_ra"]
            self.stack.attrs["cirs_ra"] = transit.attrs["cirs_ra"]

            # Copy data for first transit
            flag = (transit.weight[:] > 0.0).astype(np.int)
            if self.weight == "inverse_variance":
                coeff = transit.weight[:]
            else:
                coeff = flag.astype(np.float32)

            self.stack.beam[:] = coeff * transit.beam[:]
            self.stack.weight[:] = (coeff ** 2) * invert_no_zero(transit.weight[:])
            self.stack.nsample[:] = flag.astype(np.int)

            self.variance = coeff * np.abs(transit.beam[:]) ** 2
            self.pseudo_variance = coeff * transit.beam[:] ** 2
            self.norm = coeff

        else:
            if list(transit.beam.shape) != list(self.stack.beam.shape):
                self.log.error(
                    "Transit has different shape than stack: {}, {}".format(
                        transit.beam.shape, self.stack.beam.shape
                    )
                    + " Skipping."
                )
                return None

            self.log.info("Adding %s to stack." % transit.attrs["tag"])

            self.stack.attrs["filename"].append(transit.attrs["tag"])
            self.stack.attrs["observation_id"].append(transit.attrs["observation_id"])
            self.stack.attrs["transit_time"].append(transit.attrs["transit_time"])
            self.stack.attrs["archivefiles"] += list(transit.attrs["archivefiles"])

            # Accumulate transit data
            flag = (transit.weight[:] > 0.0).astype(np.int)
            if self.weight == "inverse_variance":
                coeff = transit.weight[:]
            else:
                coeff = flag.astype(np.float32)

            self.stack.beam[:] += coeff * transit.beam[:]
            self.stack.weight[:] += (coeff ** 2) * invert_no_zero(transit.weight[:])
            self.stack.nsample[:] += flag

            self.variance += coeff * np.abs(transit.beam[:]) ** 2
            self.pseudo_variance += coeff * transit.beam[:] ** 2
            self.norm += coeff

        return None

    def process_finish(self):
        """Normalise the stack and return the result.

        Includes the sample variance over transits within the stack.

        Returns
        -------
        stack: draco.core.containers.TrackBeam
            Stacked transits.
        """
        # Divide by norm to get average transit
        inv_norm = invert_no_zero(self.norm)
        self.stack.beam[:] *= inv_norm
        self.stack.weight[:] = invert_no_zero(self.stack.weight[:]) * self.norm ** 2

        self.variance = self.variance * inv_norm - np.abs(self.stack.beam[:]) ** 2
        self.pseudo_variance = self.pseudo_variance * inv_norm - self.stack.beam[:] ** 2

        # Calculate the covariance between the real and imaginary component
        # from the accumulated variance and psuedo-variance
        self.stack.sample_variance[0] = 0.5 * (
            self.variance + self.pseudo_variance.real
        )
        self.stack.sample_variance[1] = 0.5 * self.pseudo_variance.imag
        self.stack.sample_variance[2] = 0.5 * (
            self.variance - self.pseudo_variance.real
        )

        # Create tag
        time_range = np.percentile(self.stack.attrs["transit_time"], [0, 100])
        self.stack.attrs["tag"] = "{}_{}_to_{}".format(
            self.stack.attrs["source_name"],
            ephem.unix_to_datetime(time_range[0]).strftime("%Y%m%dT%H%M%S"),
            ephem.unix_to_datetime(time_range[1]).strftime("%Y%m%dT%H%M%S"),
        )

        return self.stack

class FilterHolographyProcessed(task.MPILoggedTask):
    """Filter holography transit DataIntervals produced by `io.QueryDatabase`
    to exclude those already processed.

    Attributes
    ----------
    processed_dir: str or list of str
        Directory or list of directories to look in for processed files.
    source: str
        The name of the holography source (as used in the holography database).
    """

    processed_dir = config.Property(
        proptype=lambda x: x if isinstance(x, list) else [x]
    )
    source = config.Property(proptype=str)

    def setup(self):
        """Get a list of existing processed files."""

        # Find processed transit files
        self.proc_transits = []
        for processed_dir in self.processed_dir:
            self.log.debug(
                "Looking for processed transits in {}...".format(processed_dir)
            )
            # Expand path
            processed_dir = path.expanduser(processed_dir)
            processed_dir = path.expandvars(processed_dir)

            try:
                processed_files = listdir(processed_dir)
            except FileNotFoundError:
                processed_files = []
            for fname in processed_files:
                if not path.splitext(fname)[1] == ".h5":
                    continue
                with ContainerBase.from_file(
                    fname, ondisk=True, distributed=False, mode="r"
                ) as fh:
                    obs_id = fh.attrs.get("observation_id", None)
                    if obs_id is not None:
                        self.proc_transits.append(obs_id)
        self.log.debug("Found {:d} processed transits.".format(len(self.proc_transits)))

        # Query database for observations of this source
        hol_obs = None
        if mpiutil.rank0:
            hol_obs = list(get_holography_obs(self.source))
        self.hol_obs = mpiutil.bcast(hol_obs, root=0)
        mpiutil.barrier()

    def next(self, intervals):
        """Filter input files to exclude those already processed.

        Parameters
        ----------
        intervals: list of chimedb.data_index.DataInterval
            List intervals to filter.

        Returns
        -------
        files: list of str
            List of files to be processed.
        """

        self.log.info("Starting next for task %s" % self.__class__.__name__)

        self.comm.Barrier()

        files = []
        for fi in intervals:
            start, end = fi[1]
            # find holography observation that overlaps this set
            this_obs = [
                o
                for o in self.hol_obs
                if (o.start_time >= start and o.start_time <= end)
                or (o.finish_time >= start and o.finish_time <= end)
                or (o.start_time <= start and o.finish_time >= end)
            ]

            if len(this_obs) == 0:
                self.log.warning(
                    "Could not find source transit in holography database for {}.".format(
                        ephem.unix_to_datetime(start)
                    )
                )
            elif this_obs[0].id in self.proc_transits:
                self.log.warning(
                    "Already processed transit for {}. Skipping.".format(
                        ephem.unix_to_datetime(start)
                    )
                )
            else:
                files += fi[0]

        self.log.info("Leaving next for task %s" % self.__class__.__name__)
        return files

def unwrap_lha(lsa, src_ra):
    """Convert LSA into HA for a source's RA. Ensures HA is monotonically increasing.

    Parameters
    ----------
    lsa: array
        Local sidereal angle.
    src_ra: float
        The RA of the source.

    Returns
    -------
    ha: array
        Hour angle.
    """
    # ensure monotonic
    start_lsa = lsa[0]
    lsa -= start_lsa
    lsa[lsa < 0] += 360.0
    lsa += start_lsa
    # subtract source RA
    return np.where(
        np.abs(lsa - src_ra) < np.abs(lsa - src_ra + 360.0),
        lsa - src_ra,
        lsa - src_ra + 360.0,
    )

def get_holography_obs(src):
    """Query database for list of all holography observations for the given
    source.

    Parameters
    ----------
    src: str
        Source name.

    Returns
    -------
    db_obs: list of ch_util.holography.HolographyObservation
        Observations of this source.
    """
    connect_database()
    db_src = holography.HolographySource.get(holography.HolographySource.name == src)
    db_obs = holography.HolographyObservation.select().where(
        holography.HolographyObservation.source == db_src
    )
    return db_obs


"""Jax Unet inference implementation."""

from concurrent.futures import thread
import functools
import math

from absl import logging
from flax import nn
import jax
from jax.interpreters import sharded_jit
import jax.numpy as jnp
import numpy as np

from REDACTED.mlperf.submissions.training.v1_0.models.unet3d.data_loading import input_reader
from REDACTED.mlperf.submissions.training.v1_0.models.unet3d.helpers import helpers
from REDACTED.mlperf.submissions.training.v1_0.models.unet3d.models import losses

def multi_device_score_fn(device_results, labels, is_real_input):
  """Returns the sum of the scores."""
  device_results = device_results.astype(jnp.bfloat16)
  labels = labels.astype(jnp.int8)
  device_results = jax.lax.psum(device_results, axis_name="cores")
  labels = jax.lax.pmax(labels, axis_name="cores")
  device_scores = losses.dice_score(device_results, labels)
  return jnp.sum(device_scores * is_real_input.astype(jnp.float32), axis=0)

def precompile_score_fn(params):
  """First sums up devices results from different cores, masks it and calls original score fn."""
  del params
  num_partitions = jax.local_device_count()
  num_total_devices = jax.device_count()
  result_partition = sharded_jit.PartitionSpec(1, num_partitions, 1, 1, 1)
  label_partition = sharded_jit.PartitionSpec(1, num_partitions, 1, 1, 1)

  s_score_fn = sharded_jit.sharded_jit(
      multi_device_score_fn,
      num_partitions=num_partitions,
      in_parts=(result_partition, label_partition, None),
      out_parts=None)
  p_score_fn = jax.pmap(
      s_score_fn,
      axis_name="cores",  # 1 replica per host.
      axis_size=num_total_devices // num_partitions,
      in_axes=(0, None, None),
      donate_argnums=(0, 1, 2))
  return p_score_fn

def reshape_for_local_replicas(np_tensor, num_local_devices):
  tensor_shape = np_tensor.shape
  new_tensor_shape = [num_local_devices, -1]
  new_tensor_shape.extend(tensor_shape[1:])
  return np.reshape(np_tensor, new_tensor_shape)

def eval_core_fn(image, crop_masks, norm_maps, model, state, norm_patch,
                 params):
  """Applies forward pass of the model and updates."""
  assert image.dtype == params["dtype"]
  assert crop_masks.dtype == jnp.int8

  assert norm_maps.dtype == params["dtype"]
  assert norm_patch.dtype == params["dtype"]
  # TODO: We might add this as a constant for TPUs.
  # norm_patch = helpers.get_norm_patch(params)
  with nn.stateful(state, mutable=False):
    model_out = model(image)
    assert model_out.dtype == params["dtype"]
    model_out_scale = model_out * norm_patch * crop_masks / norm_maps
    assert model_out_scale.dtype == params["dtype"]
  return model_out_scale

def make_eval_core_fn(params):
  """Creates an core eval fn to be called from sliding windows."""
  eval_core_function = functools.partial(eval_core_fn,
                                         params=params)
  num_partitions = params["num_eval_partitions"]
  num_eval_replicas = params["num_eval_replicas"]
  if num_partitions > 1:
    image_partition = sharded_jit.PartitionSpec(1, num_partitions,
                                                1, 1, 1)
    in_partitions = (image_partition, None, None, None, None, None)
    s_eval_core = sharded_jit.sharded_jit(eval_core_function,
                                          in_parts=in_partitions,
                                          out_parts=None)
  else:
    s_eval_core = eval_core_function
  peval_core_fn = jax.pmap(s_eval_core,
                           axis_size=num_eval_replicas,
                           in_axes=(0, 0, 0, None, None, None))

  return peval_core_fn

def merge_result_on_cpu(global_result, global_labels,
                        roi_results, roi_labels,
                        image_indices, crop_locations, is_padding, params):
  """Merges the roi region result."""
  num_local_replicas = params["num_host_eval_replicas"]
  host_eval_batch_size = params["host_eval_batch_size"]
  roi_shape = params["val_input_shape_without_channel"]
  roi_results = jax.device_get(roi_results)
  for device_id in range(num_local_replicas):
    for batch_id in range(host_eval_batch_size):

      image_id = image_indices[device_id][batch_id]
      crop_loc = crop_locations[device_id][batch_id]
      if is_padding[device_id][batch_id]:
        continue
      i, j, k, _ = crop_loc
      global_result[image_id, i:(roi_shape[0] + i), j:(roi_shape[1] + j),
                    k:(roi_shape[2] + k), :] += roi_results[device_id][batch_id]

      # Overwrite the labels.
      global_labels[image_id, i:(roi_shape[0] + i), j:(roi_shape[1] + j),
                    k:(roi_shape[2] + k)] = roi_labels[device_id][batch_id]

def prepare_eval_for_window_batches(rng, params, optimizer, state):
  """Returns precompiled eval fns."""
  del rng
  eval_core_fn_jit = make_eval_core_fn(params)

  roi_shape = params["val_input_shape"]
  eval_replica_bs = params["replica_eval_batch_size"]
  num_host_eval_replicas = params["num_host_eval_replicas"]
  eval_image_shape = [
      num_host_eval_replicas, eval_replica_bs, roi_shape[0], roi_shape[1],
      roi_shape[2], roi_shape[3]
  ]
  image = np.zeros(eval_image_shape).astype(params["dtype"])
  crop_masks = np.zeros(eval_image_shape).astype(np.int8)
  norm_maps = np.zeros(eval_image_shape).astype(params["dtype"])
  norm_patch = helpers.get_norm_patch(params).astype(params["dtype"])

  eval_core_fn_jit(image, crop_masks, norm_maps, optimizer, state, norm_patch)
  pscore_fn = precompile_score_fn(params)
  score_fn_local_device_cnt = 1
  score_batch_size = params["eval_score_fn_bs"]
  eval_result_shape = [
      score_fn_local_device_cnt, score_batch_size,
      input_reader.MAX_EVAL_IMAGE_SHAPE[0],
      input_reader.MAX_EVAL_IMAGE_SHAPE[1],
      input_reader.MAX_EVAL_IMAGE_SHAPE[2], 3
  ]
  eval_result = np.zeros(eval_result_shape).astype(jax.numpy.bfloat16)
  label_shape = [score_batch_size,
                 input_reader.MAX_EVAL_IMAGE_SHAPE[0],
                 input_reader.MAX_EVAL_IMAGE_SHAPE[1],
                 input_reader.MAX_EVAL_IMAGE_SHAPE[2],
                 input_reader.MAX_EVAL_IMAGE_SHAPE[3]]
  label = np.zeros(label_shape).astype(np.int8)

  is_real_input = np.ones([score_batch_size, 1]).astype(np.float32)

  pscore_fn(eval_result, label, is_real_input)
  return eval_core_fn_jit, pscore_fn

def evaluate_host_loop(eval_iterator, model, state, peval_core_fn, score_fn,
                       params, async_merge=True):
  """Runs eval partially on CPU (due to memory reasons) for the epoch."""
  result_merger = thread.ThreadPoolExecutor(1, "result_merger")
  score_batch_size = params["eval_score_fn_bs"]
  num_eval_samples = params["num_eval_images"]
  padding_size = num_eval_samples % score_batch_size
  num_eval_samples_with_padding = padding_size + num_eval_samples
  is_real_input = np.concatenate(
      [np.ones([num_eval_samples, 1]),
       np.zeros([padding_size, 1])]).astype(np.float32)
  device_reshape = functools.partial(
      reshape_for_local_replicas,
      num_local_devices=params["num_host_eval_replicas"])
  norm_patch = helpers.get_norm_patch(params).astype(params["dtype"])
  num_eval_steps = params["num_eval_steps"]

  global_result = np.zeros(
      shape=(num_eval_samples_with_padding,
             *input_reader.MAX_EVAL_IMAGE_SHAPE[:-1], 3),
      dtype=jnp.bfloat16)

  global_labels = np.zeros(shape=(num_eval_samples_with_padding,
                                  *input_reader.MAX_EVAL_IMAGE_SHAPE),
                           dtype=np.int8)
  for step in range(num_eval_steps):
    with jax.profiler.StepTraceAnnotation("eval", step_num=step):
      eval_ds_output = next(eval_iterator)
      eval_ds_output = jax.tree_map(lambda x: x.numpy(), eval_ds_output)
      eval_ds_output = jax.tree_map(device_reshape, eval_ds_output)
      roi_result = peval_core_fn(eval_ds_output["image"],
                                 eval_ds_output["crop_masks"],
                                 eval_ds_output["norm_maps"],
                                 model, state, norm_patch)
      merge = functools.partial(merge_result_on_cpu,
                                global_result, global_labels,
                                roi_result, eval_ds_output["label"],
                                eval_ds_output["image_idx"],
                                eval_ds_output["crop_locations"],
                                eval_ds_output["is_padding"], params)
      if async_merge:
        result_merger.submit(merge)
      else:
        merge()
  result_merger.shutdown()
  scores = []

  for i in range(0, num_eval_samples_with_padding, score_batch_size):
    gr = np.expand_dims(global_result[i: i + score_batch_size], axis=0)
    # gnm = np.expand_dims(global_norm_map[i: i + score_batch_size], axis=0)
    # gl = np.expand_dims(global_labels[i: i + score_batch_size], axis=0)
    # gnm = global_norm_map[i: i + score_batch_size]
    gl = global_labels[i: i + score_batch_size]
    score_sums = score_fn(gr, gl, is_real_input[i:i + score_batch_size])
    scores.append(score_sums)
  return scores

def update_global_result(global_result, local_result, crop_loc, params):
  """Updates the global results with window result."""
  roi_shape = params["val_input_shape"]
  num_classes = global_result.shape[-1]

  i, j, k, _ = crop_loc
  result_portion = jax.lax.dynamic_slice(
      global_result,
      [i, j, k, 0],
      [roi_shape[0], roi_shape[1], roi_shape[2], num_classes])

  local_result_without_background = local_result[:, :, :, 1:]
  local_result_without_background -= local_result[:, :, :, :1]
  global_result = jax.lax.dynamic_update_slice(
      global_result,
      result_portion + local_result_without_background,
      [i, j, k, 0])
  return global_result

def update_label(global_label, local_label, crop_loc):
  """Updates the global label with window label."""
  i, j, k, _ = crop_loc
  global_label = jax.lax.dynamic_update_slice(global_label, local_label,
                                              [i, j, k, 0])
  return global_label

def get_per_replica_eval_shapes_and_dtypes(params):
  """Returns the expected shape and dtypes for eval step inputs."""
  eval_input = {}
  roi_shape = params["val_input_shape"]
  replica_eval_batch_size = params["replica_eval_batch_size"]

  shape_5d = [
      replica_eval_batch_size, roi_shape[0], roi_shape[1], roi_shape[2],
      roi_shape[3]
  ]

  image_dtype = jnp.bfloat16 if params["use_bfloat16"] else jnp.float32
  eval_input["image"] = (shape_5d, image_dtype)
  eval_input["label"] = (shape_5d, jnp.int8)
  eval_input["crop_masks"] = (shape_5d, jnp.int8)
  eval_input["norm_maps"] = (shape_5d, image_dtype)

  eval_input["image_idx"] = ([replica_eval_batch_size], jnp.int32)
  eval_input["crop_locations"] = ([replica_eval_batch_size, 4], jnp.int32)
  eval_input["is_padding"] = ([replica_eval_batch_size], jnp.bool_)
  return eval_input

def assert_expected_shape_and_dtypes(ds_out, expected_shape_dtypes):
  """Asserts that the dataset output matches with the expected eval shapes and dtypes."""

  for data_name in expected_shape_dtypes.keys():
    logging.info("Eval shape :%s :%s :%s", data_name, ds_out[data_name].shape,
                 expected_shape_dtypes[data_name][0])
    logging.info("Eval dtype :%s :%s :%s", data_name, ds_out[data_name].dtype,
                 expected_shape_dtypes[data_name][1])
    ds_shape = list(ds_out[data_name].shape)
    expected_shape = list(expected_shape_dtypes[data_name][0])
    assert ds_shape == expected_shape
    real_dtype = jax.dtypes.canonicalize_dtype(ds_out[data_name].dtype)
    expected_dtype = jax.dtypes.canonicalize_dtype(
        expected_shape_dtypes[data_name][1])
    assert real_dtype == expected_dtype

def infeed_eval_iter(infeed_pool, eval_single_batch, infeed_devices,
                     params=None):
  """Infeeds all local devices with a single batch."""
  if params:
    # Make sure the input is as we expect.
    expected_input_config = get_per_replica_eval_shapes_and_dtypes(params)
    single_device_batch = jax.tree_map(lambda x: x[0], eval_single_batch)
    assert len(infeed_devices) == eval_single_batch["image"].shape[0]
    assert_expected_shape_and_dtypes(single_device_batch, expected_input_config)

  for index, device in enumerate(infeed_devices):
    infeed_pool.submit(
        functools.partial(device.transfer_to_infeed,
                          (eval_single_batch["image"][index],
                           eval_single_batch["label"][index],
                           eval_single_batch["crop_masks"][index],
                           eval_single_batch["norm_maps"][index],
                           eval_single_batch["image_idx"][index],
                           eval_single_batch["crop_locations"][index],
                           eval_single_batch["is_padding"][index])))

# We do not necessariliy need current_eval_step, but we will use it for
# precompile.
def evaluate_device_loop(model, state, current_eval_step, replica_id, params):
  """Runs eval for the epoch."""
  expected_input_config = get_per_replica_eval_shapes_and_dtypes(params)
  num_eval_samples = params["num_eval_images"]
  num_classes = 2
  global_results = {}
  global_labels = {}
  result_shape = (*input_reader.MAX_EVAL_IMAGE_SHAPE[:-1], num_classes)
  label_shape = (*input_reader.MAX_EVAL_IMAGE_SHAPE[:-1], 1)
  norm_patch = helpers.get_norm_patch(params).astype(params["dtype"])

  eval_core = functools.partial(eval_core_fn, params=params)
  eval_core = jax.jit(eval_core)
  update_global_result_jit = functools.partial(update_global_result,
                                               params=params)
  update_global_result_jit = jax.jit(update_global_result_jit)
  update_label_jit = jax.jit(update_label)
  for i in range(num_eval_samples):
    global_labels[i] = jnp.zeros(label_shape, dtype=jnp.int8)
    global_results[i] = jnp.zeros(result_shape, dtype=jnp.bfloat16)

  roi_shape = params["val_input_shape"]
  bs = params["replica_eval_batch_size"]

  def device_eval_loop_cond(args):
    model, state, results, labels, token, eval_step = args
    del model, state, results, labels, token
    return eval_step < params["num_eval_steps"]

  def device_eval_loop_body(args):
    model, state, results, labels, token, eval_step = args
    eval_step = eval_step + 1
    (image, label, crop_masks, norm_maps, image_idx, crop_locations,
     is_padding), token = jax.lax.infeed(
         token,
         shape=(jax.ShapedArray(expected_input_config["image"][0],
                                expected_input_config["image"][1]),
                jax.ShapedArray(expected_input_config["label"][0],
                                expected_input_config["label"][1]),
                jax.ShapedArray(expected_input_config["crop_masks"][0],
                                expected_input_config["crop_masks"][1]),
                jax.ShapedArray(expected_input_config["norm_maps"][0],
                                expected_input_config["norm_maps"][1]),
                jax.ShapedArray(expected_input_config["image_idx"][0],
                                expected_input_config["image_idx"][1]),
                jax.ShapedArray(expected_input_config["crop_locations"][0],
                                expected_input_config["crop_locations"][1]),
                jax.ShapedArray(expected_input_config["is_padding"][0],
                                expected_input_config["is_padding"][1])))
    eval_ds_output = {
        "image": image,
        "label": label,
        "crop_masks": crop_masks,
        "norm_maps": norm_maps,
        "image_idx": image_idx,
        "crop_locations": crop_locations,
        "is_padding": is_padding
    }

    # Make sure the input is as we expect.
    assert_expected_shape_and_dtypes(eval_ds_output, expected_input_config)

    batch_result = eval_core(eval_ds_output["image"],
                             eval_ds_output["crop_masks"],
                             eval_ds_output["norm_maps"], model, state,
                             norm_patch)
    batch_result = batch_result.astype(jnp.bfloat16)
    eval_ds_output["label"] = eval_ds_output["label"].astype(jnp.int8)
    assert batch_result.shape == (bs, roi_shape[0], roi_shape[1], roi_shape[2],
                                  3)
    for k in range(bs):
      local_label = eval_ds_output["label"][k]
      image_id = eval_ds_output["image_idx"][k]
      is_padding = eval_ds_output["is_padding"][k].astype(jnp.int32)
      # If padding none of the below branches should match.
      image_id = jnp.int32(image_id + is_padding * 1000)
      crop_loc = eval_ds_output["crop_locations"][k]
      single_window_result = batch_result[k]
      for i in range(params["image_index_begin"], params["image_index_end"]):
        # pylint: disable=cell-var-from-loop, g-long-lambda
        def update_while_loop_cond(args):
          (image_idx, static_image_idx, _, _, _) = args
          return image_idx == jnp.int32(static_image_idx)

        def update_result_while_loop_body(args):
          (image_idx, static_image_idx, image_result, single_window_r,
           window_loc) = args
          image_result = update_global_result_jit(image_result,
                                                  single_window_r,
                                                  window_loc)
          image_idx = -1
          return (image_idx, static_image_idx, image_result, single_window_r,
                  window_loc)

        def update_label_while_loop_body(args):
          (image_idx, static_image_idx, label_result, single_window_label,
           window_loc) = args
          label_result = update_label_jit(label_result, single_window_label,
                                          window_loc)
          image_idx = -1
          return (image_idx, static_image_idx, label_result,
                  single_window_label, window_loc)
        static_i = jnp.int32(i)
        loop_result = jax.lax.while_loop(
            update_while_loop_cond,
            update_result_while_loop_body,
            (image_id, static_i, results[i], single_window_result, crop_loc))
        results[i] = loop_result[2]

        loop_result = jax.lax.while_loop(
            update_while_loop_cond,
            update_label_while_loop_body,
            (image_id, static_i, labels[i], local_label, crop_loc))
        labels[i] = loop_result[2]

        # pylint: enable=cell-var-from-loop, g-long-lambda
    return model, state, results, labels, token, eval_step

  token = jax.lax.create_token(current_eval_step)
  (model, state, results, labels, token, _) = jax.lax.while_loop(
      device_eval_loop_cond, device_eval_loop_body,
      (model, state, global_results, global_labels, token, current_eval_step))

  for i in range(params["image_index_begin"], params["image_index_end"]):
    results[i] = jax.lax.psum(results[i], axis_name="cores")
    labels[i] = jax.lax.pmax(labels[i], axis_name="cores")
  scores = []
  score_fn = jax.jit(losses.dice_score_withoout_background)

  logging.info("num_eval_replicas:%s num_eval_images:%s",
               params["num_eval_replicas"], params["num_eval_images"])
  if params["num_eval_replicas"] < params["num_eval_images"]:
    # If we have less replicas than the number of images, simply all replicas
    # will compute the all scores.
    for i in range(params["image_index_begin"], params["image_index_end"]):
      results[i] = jnp.expand_dims(results[i], axis=0)
      labels[i] = jnp.expand_dims(labels[i], axis=0)
      scores.append(score_fn(results[i], labels[i]))
    return scores
  # This is a function that will never be executed, via jax.lax.cond dead
  # branch. We add this, to make sure that all psums are executed before
  # computing dice_scores. Otherwise, the compiler may try to overlap the
  # psum computation with score computation, which is done by one core at a
  # time.
  def force_execution_of_all_psums():
    r = []
    for i in range(params["image_index_begin"], params["image_index_end"]):
      result = results[i]
      label = labels[i]
      result = jnp.expand_dims(result, axis=0)
      label = jnp.expand_dims(label, axis=0)
      r.append(score_fn(result, label))
    return r
  dummy_fn = jax.jit(force_execution_of_all_psums)

  def score_loop_body(x):
    result, label = x
    result = jnp.expand_dims(result, axis=0)
    label = jnp.expand_dims(label, axis=0)
    return score_fn(result, label)
  score_loop_body = jax.jit(score_loop_body)

  def compute_scores():
    r = []
    for image_idx in range(params["image_index_begin"],
                           params["image_index_end"]):
      r.append(jax.lax.cond(
          replica_id == image_idx,
          lambda x: score_loop_body(x),  # pylint: disable=unnecessary-lambda
          lambda x: jnp.zeros([1, 2]).astype(jnp.float32),
          (results[image_idx], labels[image_idx])))
    return r
  compute_scores = jax.jit(compute_scores)
  args = (results, labels)
  scores = jax.lax.cond(
      replica_id == -1,  # replica id will never be -1
      lambda x: dummy_fn(),  # pylint: disable=unnecessary-lambda
      lambda x: compute_scores(),
      args)

  scores = [jax.lax.psum(r, axis_name="cores") for r in scores]
  return scores

def get_eval_fn(rng, optimizer, state, params, infeed_pool):
  """Returns evaluate fn."""
  if params["use_eval_device_loop"]:

    per_eval_image_size = int(
        math.ceil(params["num_eval_images"] / params["num_eval_passes"]))
    eval_fns = []
    replica_ids = np.arange(jax.local_device_count() * jax.host_id(),
                            jax.local_device_count() *
                            (jax.host_id() + 1)).astype(np.int32)

    for i in range(0, params["num_eval_images"], per_eval_image_size):
      params["image_index_begin"] = i
      params["image_index_end"] = min(i + per_eval_image_size,
                                      params["num_eval_images"])
      logging.info("Compiling eval for begin:%d, end:%d",
                   params["image_index_begin"], params["image_index_end"])
      eval_device_loop = functools.partial(evaluate_device_loop, params=params)
      p_eval_device_loop = jax.pmap(
          eval_device_loop, in_axes=(None, None, 0, 0), axis_name="cores")
      precompile_steps = (np.zeros([jax.local_device_count()], dtype=np.int64) +
                          np.int64(params["num_eval_steps"])).astype(np.int32)
      # This should just compile and not execute.
      p_eval_device_loop(optimizer, state, precompile_steps, replica_ids)
      eval_fns.append(p_eval_device_loop)

    infeed_devices = jax.local_devices()
    device_reshape = functools.partial(
        reshape_for_local_replicas,
        num_local_devices=params["num_host_eval_replicas"])
    def evaluate_with_infeed(eval_iterator, model, state):
      scores = []
      for i in range(params["num_eval_passes"]):
        current_step = np.zeros([jax.local_device_count()], dtype=np.int32)
        p_eval_device_loop = eval_fns[i]
        score = p_eval_device_loop(model, state, current_step, replica_ids)
        for s in range(params["num_eval_steps"]):
          with jax.profiler.StepTraceAnnotation("eval", step_num=s):
            eval_ds_output = next(eval_iterator)
            eval_ds_output = jax.tree_map(lambda x: x.numpy(), eval_ds_output)
            eval_single_batch = jax.tree_map(device_reshape, eval_ds_output)
            infeed_eval_iter(infeed_pool, eval_single_batch, infeed_devices)
        scores.extend(score)
      return scores
    return evaluate_with_infeed
  else:
    eval_core_fn_jit, pscore_fn = prepare_eval_for_window_batches(rng, params,
                                                                  optimizer,
                                                                  state)
    evaluate_fn = functools.partial(evaluate_host_loop,
                                    peval_core_fn=eval_core_fn_jit,
                                    score_fn=pscore_fn,
                                    params=params,
                                    async_merge=True)
  return evaluate_fn