from data.voc0712 import *
from utils.augmentations import SSDAugmentation
from layers.modules import MultiBoxLoss
from ssd import build_ssd
import os
import sys
import time
import torch
from torch.autograd import Variable
import torch.nn as nn
import torch.optim as optim
import torch.backends.cudnn as cudnn
import torch.nn.init as init
import torch.utils.data as data
import numpy as np
import argparse
from data import *


def str2bool(v):
    return v.lower() in ("yes", "true", "t", "1")


parser = argparse.ArgumentParser(
    description='Single Shot MultiBox Detector Training With Pytorch')
train_set = parser.add_mutually_exclusive_group()
parser.add_argument('--dataset', default='VOC', choices=['VOC', 'COCO'],
                    type=str, help='VOC or COCO')
parser.add_argument('--dataset_root', default=VOC_ROOT,
                    help='Dataset root directory path')
parser.add_argument('--basenet', default='vgg16_reducedfc.pth',
                    help='Pretrained base model')
parser.add_argument('--batch_size', default=32, type=int,
                    help='Batch size for training')
parser.add_argument('--resume', default=None, type=str,
                    help='Checkpoint state_dict file to resume training from')
parser.add_argument('--start_iter', default=0, type=int,
                    help='Resume training at this iter')
parser.add_argument('--num_workers', default=4, type=int,
                    help='Number of workers used in dataloading')
parser.add_argument('--cuda', default=True, type=str2bool,
                    help='Use CUDA to train model')
parser.add_argument('--lr', '--learning-rate', default=1e-3, type=float,
                    help='initial learning rate')
parser.add_argument('--momentum', default=0.9, type=float,
                    help='Momentum value for optim')
parser.add_argument('--weight_decay', default=5e-4, type=float,
                    help='Weight decay for SGD')
parser.add_argument('--gamma', default=0.1, type=float,
                    help='Gamma update for SGD')
parser.add_argument('--visdom', default=True, type=str2bool,
                    help='Use visdom for loss visualization')
parser.add_argument('--save_folder', default='weights/',
                    help='Directory for saving checkpoint models')
args = parser.parse_args()


if torch.cuda.is_available():
    if args.cuda:
        torch.set_default_tensor_type('torch.cuda.FloatTensor')
    if not args.cuda:
        print("WARNING: It looks like you have a CUDA device, but aren't " +
              "using CUDA.\nRun with --cuda for optimal training speed.")
        torch.set_default_tensor_type('torch.FloatTensor')
else:
    torch.set_default_tensor_type('torch.FloatTensor')

if not os.path.exists(args.save_folder):
    os.mkdir(args.save_folder)


def train():
    if args.dataset == 'COCO':
        if args.dataset_root == VOC_ROOT:
            if not os.path.exists(COCO_ROOT):
                parser.error('Must specify dataset_root if specifying dataset')
            print("WARNING: Using default COCO dataset_root because " +
                  "--dataset_root was not specified.")
            args.dataset_root = COCO_ROOT
        cfg = coco
        dataset = COCODetection(root=args.dataset_root,
                                transform=SSDAugmentation(cfg['min_dim'],
                                                          MEANS))
    elif args.dataset == 'VOC':
        if args.dataset_root == COCO_ROOT:
            parser.error('Must specify dataset if specifying dataset_root')
        cfg = voc
        dataset = VOCDetection(root=args.dataset_root,
                               transform=SSDAugmentation(cfg['min_dim'],
                                                         MEANS))

    if args.visdom:
        import visdom
        viz = visdom.Visdom()

    # 构建一个SSD网络
    ssd_net = build_ssd('train', cfg['min_dim'], cfg['num_classes'])
    net = ssd_net

    if args.cuda:
        net = torch.nn.DataParallel(ssd_net)
        cudnn.benchmark = True

    if args.resume:
        print('Resuming training, loading {}...'.format(args.resume))
        ssd_net.load_weights(args.resume)
    else:
        vgg_weights = torch.load(args.save_folder + args.basenet)
        print('Loading base network...')
        ssd_net.vgg.load_state_dict(vgg_weights)

    if args.cuda:
        net = net.cuda()

    if not args.resume:
        print('Initializing weights...')
        # initialize newly added layers' weights with xavier method
        ssd_net.extras.apply(weights_init)
        ssd_net.loc.apply(weights_init)
        ssd_net.conf.apply(weights_init)

    optimizer = optim.SGD(net.parameters(), lr=args.lr, momentum=args.momentum,
                          weight_decay=args.weight_decay)
    criterion = MultiBoxLoss(cfg['num_classes'], 0.5, True, 0, True, 3, 0.5,
                             False, args.cuda)

    net.train()
    # loss counters
    loc_loss = 0
    conf_loss = 0
    epoch = 0
    print('Loading the dataset...')

    epoch_size = len(dataset) // args.batch_size
    print('Training SSD on:', dataset.name)
    print('Using the specified args:')
    print(args)

    step_index = 0

    if args.visdom:
        vis_title = 'SSD.PyTorch on ' + dataset.name
        vis_legend = ['Loc Loss', 'Conf Loss', 'Total Loss']
        iter_plot = create_vis_plot('Iteration', 'Loss', vis_title, vis_legend)
        epoch_plot = create_vis_plot('Epoch', 'Loss', vis_title, vis_legend)

    data_loader = data.DataLoader(dataset, args.batch_size,
                                  num_workers=args.num_workers,
                                  shuffle=True, collate_fn=detection_collate,
                                  pin_memory=True)
    # create batch iterator
    batch_iterator = iter(data_loader)
    for iteration in range(args.start_iter, cfg['max_iter']):
        if args.visdom and iteration != 0 and (iteration % epoch_size == 0):
            update_vis_plot(epoch, loc_loss, conf_loss, epoch_plot, None,
                            'append', epoch_size)
            # reset epoch loss counters
            loc_loss = 0
            conf_loss = 0
            epoch += 1

        if iteration in cfg['lr_steps']:
            step_index += 1
            adjust_learning_rate(optimizer, args.gamma, step_index)

        # load train data
        images, targets = next(batch_iterator)

        if args.cuda:
            images = Variable(images.cuda())
            targets = [Variable(ann.cuda(), volatile=True) for ann in targets]
        else:
            images = Variable(images)
            targets = [Variable(ann, volatile=True) for ann in targets]
        # forward
        t0 = time.time()
        out = net(images)
        # backprop
        optimizer.zero_grad()
        loss_l, loss_c = criterion(out, targets)
        loss = loss_l + loss_c
        loss.backward()
        optimizer.step()
        t1 = time.time()
        loc_loss += loss_l.data[0]
        conf_loss += loss_c.data[0]

        if iteration % 10 == 0:
            print('timer: %.4f sec.' % (t1 - t0))
            print('iter ' + repr(iteration) + ' || Loss: %.4f ||' % (loss.data[0]), end=' ')

        if args.visdom:
            update_vis_plot(iteration, loss_l.data[0], loss_c.data[0],
                            iter_plot, epoch_plot, 'append')

        if iteration != 0 and iteration % 5000 == 0:
            print('Saving state, iter:', iteration)
            torch.save(ssd_net.state_dict(), 'weights/ssd300_COCO_' +
                       repr(iteration) + '.pth')
    torch.save(ssd_net.state_dict(),
               args.save_folder + '' + args.dataset + '.pth')


def adjust_learning_rate(optimizer, gamma, step):
    """Sets the learning rate to the initial LR decayed by 10 at every
        specified step
    # Adapted from PyTorch Imagenet example:
    # https://github.com/pytorch/examples/blob/master/imagenet/main.py
    """
    lr = args.lr * (gamma ** (step))
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr


def xavier(param):
    init.xavier_uniform(param)


def weights_init(m):
    if isinstance(m, nn.Conv2d):
        xavier(m.weight.data)
        m.bias.data.zero_()


def create_vis_plot(_xlabel, _ylabel, _title, _legend):
    return viz.line(
        X=torch.zeros((1,)).cpu(),
        Y=torch.zeros((1, 3)).cpu(),
        opts=dict(
            xlabel=_xlabel,
            ylabel=_ylabel,
            title=_title,
            legend=_legend
        )
    )


def update_vis_plot(iteration, loc, conf, window1, window2, update_type,
                    epoch_size=1):
    viz.line(
        X=torch.ones((1, 3)).cpu() * iteration,
        Y=torch.Tensor([loc, conf, loc + conf]).unsqueeze(0).cpu() / epoch_size,
        win=window1,
        update=update_type
    )
    # initialize epoch plot on first iteration
    if iteration == 0:
        viz.line(
            X=torch.zeros((1, 3)).cpu(),
            Y=torch.Tensor([loc, conf, loc + conf]).unsqueeze(0).cpu(),
            win=window2,
            update=True
        )


if __name__ == '__main__':
    train()

def build_ssd(phase, size=300, num_classes=21):
    if phase != "test" and phase != "train":
        print("ERROR: Phase: " + phase + " not recognized")
        return
    if size != 300:
        print("ERROR: You specified size " + repr(size) + ". However, " +
              "currently only SSD300 (size=300) is supported!")
        return

    base_, extras_, head_ = multibox(vgg(base[str(size)], 3),
                     add_extras(extras[str(size)], 1024),
                     mbox[str(size)],
                     num_classes)
    return SSD(phase, size, base_, extras_, head_, num_classes)

base_, extras_, head_ = multibox(vgg(base[str(size)], 3),
				 add_extras(extras[str(size)], 1024),
				 mbox[str(size)],
				 num_classes)

# This function is derived from torchvision VGG make_layers()
# https://github.com/pytorch/vision/blob/master/torchvision/models/vgg.py
def vgg(cfg, i, batch_norm=False):
    # cfg 是  [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'C', 512, 512, 512, 'M', 512, 512, 512],
    layers = []
    in_channels = i
    for v in cfg:
        if v == 'M':
            layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
        # When ceil_mode True, will use ceil instead of floor to compute the output shape
        elif v == 'C':
            layers += [nn.MaxPool2d(kernel_size=2, stride=2, ceil_mode=True)]
        else:
            conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1)
            if batch_norm:
                layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)]
            else:
                layers += [conv2d, nn.ReLU(inplace=True)]
            in_channels = v

    # 参见论文第3部分开头，
    # change pool5 from 2*2-s2 to 3*3-s1
    # convert vgg fc6 & fc7 to conv layers
    # and use atrous algorithm to fill the "hole"
    pool5 = nn.MaxPool2d(kernel_size=3, stride=1, padding=1)
    conv6 = nn.Conv2d(512, 1024, kernel_size=3, padding=6, dilation=6)
    conv7 = nn.Conv2d(1024, 1024, kernel_size=1)
    layers += [pool5, conv6,
               nn.ReLU(inplace=True), conv7, nn.ReLU(inplace=True)]
    return layers

def add_extras(cfg, i, batch_norm=False):
    # Extra layers added to VGG for feature scaling
    layers = []
    in_channels = i
    flag = False  # 是否用3*3 conv的标志
    for k, v in enumerate(cfg):
        if in_channels != 'S':
            if v == 'S':
                layers += [nn.Conv2d(in_channels, cfg[k + 1], kernel_size=(1, 3)[flag], stride=2, padding=1)]
            else:
                layers += [nn.Conv2d(in_channels, v, kernel_size=(1, 3)[flag])]
            flag = not flag
        in_channels = v
    return layers

def multibox(vgg, extra_layers, cfg, num_classes):
    loc_layers = []
    conf_layers = []
    vgg_source = [21, -2]
    for k, v in enumerate(vgg_source):
	    loc_layers += [nn.Conv2d(vgg[v].out_channels,
				     cfg[k] * 4, kernel_size=3, padding=1)]
	    conf_layers += [nn.Conv2d(vgg[v].out_channels,
				      cfg[k] * num_classes, kernel_size=3, padding=1)]
    for k, v in enumerate(extra_layers[1::2], 2):
	    loc_layers += [nn.Conv2d(v.out_channels, cfg[k]
				     * 4, kernel_size=3, padding=1)]
	    conf_layers += [nn.Conv2d(v.out_channels, cfg[k]
				      * num_classes, kernel_size=3, padding=1)]
    return vgg, extra_layers, (loc_layers, conf_layers)

from __future__ import division
from math import sqrt as sqrt
from itertools import product as product
import torch


class PriorBox(object):
	"""Compute priorbox coordinates in center-offset form for each source
	feature map.
	"""

	def __init__(self, cfg):
		super(PriorBox, self).__init__()
		self.image_size = cfg['min_dim']
		# number of priors for feature map location (either 4 or 6)
		self.num_priors = len(cfg['aspect_ratios'])
		self.variance = cfg['variance'] or [0.1]
		self.feature_maps = cfg['feature_maps']
		self.min_sizes = cfg['min_sizes']
		self.max_sizes = cfg['max_sizes']
		self.steps = cfg['steps']
		self.aspect_ratios = cfg['aspect_ratios']
		self.clip = cfg['clip']
		self.version = cfg['name']
		for v in self.variance:
			if v <= 0:
				raise ValueError('Variances must be greater than 0')

	def forward(self):
		mean = []
		for k, f in enumerate(self.feature_maps):  # 存放的是feature map的尺寸:38,19,10,5,3,1
			# steps=[8,16,32,64,100,300]. f_k大约为feature map的尺寸
			f_k = self.image_size / self.steps[k]

			# 求解每个feature map上的每个位置的中心
			for i, j in product(range(f), repeat=2):

				# unit center x,y
				cx = (j + 0.5) / f_k
				cy = (i + 0.5) / f_k

				# aspect_ratio: 1
				# rel size: min_size
				s_k = self.min_sizes[k] / self.image_size
				mean += [cx, cy, s_k, s_k]
				# aspect_ratio: 1
				# rel size: sqrt(s_k * s_(k+1))
				s_k_prime = sqrt(s_k * (self.max_sizes[k] / self.image_size))
				mean += [cx, cy, s_k_prime, s_k_prime]

				# rest of aspect ratios
				for ar in self.aspect_ratios[k]:
					mean += [cx, cy, s_k * sqrt(ar), s_k / sqrt(ar)]
					mean += [cx, cy, s_k / sqrt(ar), s_k * sqrt(ar)]

		# 执行完上述操作后，总共有 38*38*4 + 19*19*6 + 10*10*6 + 5*5*6 + 3*3*4 + 1*1*4 = 8732个default box
		# back to torch land
		output = torch.Tensor(mean).view(-1, 4)
		# print(output[5772:5776])

		# 将输入input张量每个元素的夹紧到区间 [min,max]之内
		if self.clip:
			output.clamp_(max=1, min=0)
		return output

optimizer = optim.SGD(net.parameters(), lr=args.lr, momentum=args.momentum,
                          weight_decay=args.weight_decay)
    criterion = MultiBoxLoss(cfg['num_classes'], 0.5, True, 0, True, 3, 0.5,
                             False, args.cuda)

# -*- coding: utf-8 -*-
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
from data import coco as cfg
from ..box_utils import match, log_sum_exp


class MultiBoxLoss(nn.Module):
	"""SSD Weighted Loss Function
	Compute Targets:
		1) Produce Confidence Target Indices by matching  ground truth boxes
		   with (default) 'priorboxes' that have jaccard index > threshold parameter
		   (default threshold: 0.5).
		2) Produce localization target by 'encoding' variance into offsets of ground
		   truth boxes and their matched  'priorboxes'.
		3) Hard negative mining to filter the excessive number of negative examples
		   that comes with using a large number of default bounding boxes.
		   (default negative:positive ratio 3:1)
	Objective Loss:
		L(x,c,l,g) = (Lconf(x, c) + αLloc(x,l,g)) / N
		Where, Lconf is the CrossEntropy Loss and Lloc is the SmoothL1 Loss
		weighted by α which is set to 1 by cross val.
		Args:
			c: class confidences,
			l: predicted boxes,
			g: ground truth boxes
			N: number of matched default boxes
		See: https://arxiv.org/pdf/1512.02325.pdf for more details.
	"""

	def __init__(self, num_classes, overlap_thresh, prior_for_matching,
				 bkg_label, neg_mining, neg_pos, neg_overlap, encode_target,
				 use_gpu=True):
		super(MultiBoxLoss, self).__init__()
		self.use_gpu = use_gpu
		self.num_classes = num_classes
		self.threshold = overlap_thresh  # iou阈值
		self.background_label = bkg_label  # 背景标签 0
		self.encode_target = encode_target
		self.use_prior_for_matching = prior_for_matching
		self.do_neg_mining = neg_mining
		self.negpos_ratio = neg_pos  # 负正样例比
		self.neg_overlap = neg_overlap  # 负样本阈值 没用到
		self.variance = cfg['variance']

	def forward(self, predictions, targets):
		"""Multibox Loss
		Args:
			predictions (tuple): A tuple containing loc preds, conf preds,
			and prior boxes from SSD net.
				conf shape: torch.size(batch_size,num_priors,num_classes)
				loc shape: torch.size(batch_size,num_priors,4)
				priors shape: torch.size(num_priors,4)

			targets (tensor): Ground truth boxes and labels for a batch,
				shape: [batch_size,num_objs,5] (last idx is the label).
		"""
		# loc_data [batch, 8732, 4]
		# conf_data [batch, 8732, 21]
		# priors  [8732, 4]
		loc_data, conf_data, priors = predictions
		num = loc_data.size(0)
		priors = priors[:loc_data.size(1), :]  # 维度不变
		num_priors = (priors.size(0))  # 8732
		num_classes = self.num_classes  # 21

		# match priors (default boxes) and ground truth boxes
		loc_t = torch.Tensor(num, num_priors, 4)
		conf_t = torch.LongTensor(num, num_priors)
		for idx in range(num):
			# target shape  [batch_size,num_objs,5]
			truths = targets[idx][:, :-1].data  # [num_objs,4],即坐标位置
			labels = targets[idx][:, -1].data  # 最后一维是类别的label,因为用的是-1，不是-1:,所以返回的shape是 [num_objs]
			defaults = priors.data

			# ！！！关键函数！！！ default box和gt box之前的匹配
			#
			# 这里输入的defaults是基于中心表示的，输入的truths是基于左上角和右下角坐标表示的，在match函数中会有相应的处理
			#
			# 最后两个个参数loc_t, conf_t 其实是输出参数，
			# loc_t在调用后，返回的是第idx图片张输入时，各个default box 匹配到的gt box，相对于default box的偏移量
			# conf_t在调用后，返回的是第idx张图片输入时，各个default box匹配到的gt box的类别label值。
			# 注意，不一定所有的default box都有label, 例如有的default box 和 匹配到的gt box overlap太低，就被视作了背景
			match(self.threshold, truths, defaults, self.variance, labels,
				  loc_t, conf_t, idx)
		if self.use_gpu:
			loc_t = loc_t.cuda()
			conf_t = conf_t.cuda()
		# wrap targets
		# # 用Variable封装loc_t, 新版本的 PyTorch 无需这么做, 只需要将 requires_grad 属性设置为 True 就行了
		# loc_t: shape [batch,num_priors,4] encoded offsets to learn
		# conf_t: shape [batch,num_priors] top class label for each prior
		loc_t = Variable(loc_t, requires_grad=False)
		conf_t = Variable(conf_t, requires_grad=False)

		# 筛选出那些 label不是0的default box.类型为bool pos: shape[batch,num_priors]
		pos = conf_t > 0  # 即正样本下标
		# 统计label不是0的default box数量， 所有的default box数量为8732 [batch,1]
		num_pos = pos.sum(dim=1, keepdim=True)

		# Localization Loss (Smooth L1)
		# Shape: [batch,num_priors,4]
		# pos_idx: [batch, num_priors, 4]，这里的pos_idx作用就是论文（2）式中的X_{i,j}^{p}
		pos_idx = pos.unsqueeze(pos.dim()).expand_as(loc_data)
		loc_p = loc_data[pos_idx].view(-1, 4)  # 预测的offset
		loc_t = loc_t[pos_idx].view(-1, 4)  # 匹配的gt的offset
		# smooth-L1 loss
		loss_l = F.smooth_l1_loss(loc_p, loc_t, size_average=False)

		# Compute max conf across batch for hard negative mining
		# 计算最大置信度，以进行难负样本挖掘
		# conf_data: [batch,num_priors,num_classes]
		batch_conf = conf_data.view(-1, self.num_classes)  # batch_conf: [batch * num_priors,num_classes]

		# batch_conf.gather(1, conf_t.view(-1, 1)) 返回的是每个default box 真实label 对应的预测置信度
		# batch_conf.gather(1, conf_t.view(-1, 1)) shape 为 [batch*num_priors,1]
		# 下式求的是论文（3）式
		loss_c = log_sum_exp(batch_conf) - batch_conf.gather(1, conf_t.view(-1, 1))

		# Hard Negative Mining
		# 难负样本挖掘, 按照loss进行排序, 取loss最大的负样本参与更新
		# loss_c shape 为[batch*num_priors,1]

		# 先让那些正样本，也就是default box 对应的 gt lable 不是背景的样本，先设置为0
		loss_c[pos] = 0  # filter out pos boxes for now
		# loss_c 还原回到 [batch, num_priors]
		loss_c = loss_c.view(num, -1)
		# 对loss_c的下标进行降序排序，得到排序后的下标
		# 因为我们在前面将正样本的loss_c置为了0，所以排在前面的都是那些负样本的loss_c
		_, loss_idx = loss_c.sort(1, descending=True)
		# 对下标进行排序，得到下标的下标 ，shape 为[batch,num_priors]
		_, idx_rank = loss_idx.sort(1)
		# pos的shape是[batch,num_priors]， num_pos统计的是每个样本的正样本数目，也就是obj的数量,shape 为[batch,1]
		num_pos = pos.long().sum(1, keepdim=True)
		# 根据正样本的数量，确定负样本的数量
		num_neg = torch.clamp(self.negpos_ratio * num_pos, max=pos.size(1) - 1)
		# 获取到负样本的下标类型为bool，shape[batch,num_priors]
		neg = idx_rank < num_neg.expand_as(idx_rank)

		# Confidence Loss Including Positive and Negative Examples
		# conf_data的shape是[batch,num_priors,num_classes]
		# pos的shape是 [batch,num_priors],类型为bool，标定某个defaul box是否是正样本
		# pos_idx 的shape扩张后变为 [batch,num_priors，num_classes]，也是bool类型
		pos_idx = pos.unsqueeze(2).expand_as(conf_data)
		# neg 同理 pos
		neg_idx = neg.unsqueeze(2).expand_as(conf_data)
		# 只要样本是正样本或者负样本，就选出来，准备算到总的loss里面
		# conf_p的shape 是 [batch*num_priors,num_classes]
		conf_p = conf_data[(pos_idx + neg_idx).gt(0)].view(-1, self.num_classes)

		# 只要是正或者负样本，就取出来，targets_weighted表示它的真实标签
		# targets_weighted 的shape为[batch,num_priors]
		targets_weighted = conf_t[(pos + neg).gt(0)]
		loss_c = F.cross_entropy(conf_p, targets_weighted, size_average=False)

		# Sum of losses: L(x,c,l,g) = (Lconf(x, c) + αLloc(x,l,g)) / N

		N = num_pos.data.sum()
		loss_l /= N
		loss_c /= N
		return loss_l, loss_c

def match(threshold, truths, priors, variances, labels, loc_t, conf_t, idx):
	"""Match each prior box with the ground truth box of the highest jaccard
	overlap, encode the bounding boxes, then return the matched indices
	corresponding to both confidence and location preds.
	Args:
		threshold: (float) The overlap threshold used when mathing boxes.
		truths: (tensor) Ground truth boxes, Shape: [num_obj, 4].
		priors: (tensor) Prior boxes from priorbox layers, Shape: [n_priors,4].
		variances: (tensor) Variances corresponding to each prior coord,
			Shape: [num_priors, 4].
		labels: (tensor) All the class labels for the image, Shape: [num_obj].

		loc_t: (tensor) Tensor to be filled w/ endcoded location targets.
		conf_t: (tensor) Tensor to be filled w/ matched indices for conf preds.
		idx: (int) current batch index

	Return:
		The matched indices corresponding to 1)location and 2)confidence preds.
	"""
	# jaccard index , shape [num_objects,n_priors]
	overlaps = jaccard(
		truths,
		point_form(priors)
	)

	# 在论文的Matching Strategy说到，不仅对每个GT box,去match最匹配的default box，也对每个default box，去match最匹配的GT-box
	# 这样一来，可以让网络预测出更多的overlapping的default box，而不要让网络去选择一个最好的，至于最后怎么选择，用non-max suppression即可。

	# (Bipartite Matching)
	# [num_objects，1]
	# 表示和每个GT box重叠得最好的那个default box
	best_prior_overlap, best_prior_idx = overlaps.max(1, keepdim=True)
	# [1,num_priors]
	# 表示和每个default box重叠得最好的那个GT box
	best_truth_overlap, best_truth_idx = overlaps.max(0, keepdim=True)

	# shape [num_objects]
	best_truth_idx.squeeze_(0)
	best_truth_overlap.squeeze_(0)

	# shape [num_prior]
	best_prior_idx.squeeze_(1)
	best_prior_overlap.squeeze_(1)

	# 该语句会将与gt box匹配度最好的prior box 的交并比置为 2, 确保其最大, 以免防止某些 gtbox 没有匹配的 priorbox.
	best_truth_overlap.index_fill_(0, best_prior_idx, 2)  # ensure best prior
	# TODO refactor: index  best_prior_idx with long tensor

	# ensure every gt matches with its prior of max overlap
	for j in range(best_prior_idx.size(0)):
		best_truth_idx[best_prior_idx[j]] = j

	# best_truth_idx表示的是每个default box 对应的最佳匹配的 gt box下标，下标元素的取值范围是 0~(num_objs-1)
	# 元素的个数，也就是len(best_truth_idx)=num_priors, 也就是defalut box的个数。
	#
	# truth又是一个shape为[num_objs,4]的tensor，表示的是这么多个objs的box
	#
	# 所以truths[best_truth_idx] 代表的就是 每个default box 对应的最匹配的 gt box
	matches = truths[best_truth_idx]  # Shape: [num_priors,4]
	# conf是每个default box对应最匹配的gt box的类别，记住，0代表背景
	conf = labels[best_truth_idx] + 1  # Shape: [num_priors]

	# 对于每个default box，如果它与某个gt box的overlap小于阈值，那么就将这个default box的类别设置为0，即背景
	conf[best_truth_overlap < threshold] = 0  # label as background
	""" 把每个defalut box 匹配到的 gt box转化为 中心坐标的形式 """
	loc = encode(matches, priors, variances)
	loc_t[idx] = loc  # [num_priors,4] encoded offsets to learn
	conf_t[idx] = conf  # [num_priors] top class label for each prior

from __future__ import print_function
import sys
import os
import argparse
import torch
import torch.nn as nn
import torch.backends.cudnn as cudnn
import torchvision.transforms as transforms
from torch.autograd import Variable
from data import VOC_ROOT, VOC_CLASSES as labelmap
from PIL import Image
from data import VOCAnnotationTransform, VOCDetection, BaseTransform, VOC_CLASSES
import torch.utils.data as data
from ssd import build_ssd

parser = argparse.ArgumentParser(description='Single Shot MultiBox Detection')
parser.add_argument('--trained_model', default='weights/ssd_300_VOC0712.pth',
                    type=str, help='Trained state_dict file path to open')
parser.add_argument('--save_folder', default='eval/', type=str,
                    help='Dir to save results')
parser.add_argument('--visual_threshold', default=0.6, type=float,
                    help='Final confidence threshold')
parser.add_argument('--cuda', default=True, type=bool,
                    help='Use cuda to train model')
parser.add_argument('--voc_root', default=VOC_ROOT, help='Location of VOC root directory')
parser.add_argument('-f', default=None, type=str, help="Dummy arg so we can load in Jupyter Notebooks")
args = parser.parse_args()

if args.cuda and torch.cuda.is_available():
    torch.set_default_tensor_type('torch.cuda.FloatTensor')
else:
    torch.set_default_tensor_type('torch.FloatTensor')

if not os.path.exists(args.save_folder):
    os.mkdir(args.save_folder)


def test_net(save_folder, net, cuda, testset, transform, thresh):
    # dump predictions and assoc. ground truth to text file for now
    filename = save_folder+'test1.txt'
    num_images = len(testset)
    for i in range(num_images):
        print('Testing image {:d}/{:d}....'.format(i+1, num_images))
        img = testset.pull_image(i)
        img_id, annotation = testset.pull_anno(i)
        x = torch.from_numpy(transform(img)[0]).permute(2, 0, 1)
        x = Variable(x.unsqueeze(0))

        with open(filename, mode='a') as f:
            f.write('\nGROUND TRUTH FOR: '+img_id+'\n')
            for box in annotation:
                f.write('label: '+' || '.join(str(b) for b in box)+'\n')
        if cuda:
            x = x.cuda()

        y = net(x)      # forward pass
        detections = y.data
        # scale each detection back up to the image
        scale = torch.Tensor([img.shape[1], img.shape[0],
                             img.shape[1], img.shape[0]])
        pred_num = 0
        for i in range(detections.size(1)):
            j = 0
            while detections[0, i, j, 0] >= 0.6:
                if pred_num == 0:
                    with open(filename, mode='a') as f:
                        f.write('PREDICTIONS: '+'\n')
                score = detections[0, i, j, 0]
                label_name = labelmap[i-1]
                pt = (detections[0, i, j, 1:]*scale).cpu().numpy()
                coords = (pt[0], pt[1], pt[2], pt[3])
                pred_num += 1
                with open(filename, mode='a') as f:
                    f.write(str(pred_num)+' label: '+label_name+' score: ' +
                            str(score) + ' '+' || '.join(str(c) for c in coords) + '\n')
                j += 1


def test_voc():
    # load net
    num_classes = len(VOC_CLASSES) + 1 # +1 background
    net = build_ssd('test', 300, num_classes) # initialize SSD
    net.load_state_dict(torch.load(args.trained_model))
    net.eval()  # 设置是eval()模式，所以不会更新参数
    print('Finished loading model!')
    # load data
    testset = VOCDetection(args.voc_root, [('2007', 'test')], None, VOCAnnotationTransform())
    if args.cuda:
        net = net.cuda()
        cudnn.benchmark = True
    # evaluation
    test_net(args.save_folder, net, args.cuda, testset,
             BaseTransform(net.size, (104, 117, 123)),
             thresh=args.visual_threshold)

if __name__ == '__main__':
    test_voc()

File "XXX/SSD/layers/modules/multibox_loss.py", line 127, in forward    
    loss_c[pos] = 0  # filter out pos boxes for now
IndexError: 
    The shape of the mask [32, 8732] at index 0does not match the shape
    of the indexed tensor [279424,1] at index 0

# loss_c 还原回到 [batch, num_priors]
loss_c = loss_c.view(num, -1)
# 先让那些正样本，也就是default box 对应的 gt lable 不是背景的样本，先设置为0
loss_c[pos] = 0  # filter out pos boxes for now

# issue #173, reference https://github.com/amdegroot/ssd.pytorch/issues/173
N = num_pos.data.sum().double()
loss_l=loss_l.double()
loss_c = loss_c.double()
loss_l /= N
loss_c /= N

IndexError: invalid index of a 0-dim tensor. Use tensor.item() to convert 
a 0-dim tensor to a Python number

UserWarning: volatile was removed and now has no effect. 
    Use with `torch.no_grad()`: instead

with torch.no_grad():
    self.priors = Variable(self.priorbox.forward())

userwarning: size_average and reduce args will be deprecated, 
    please use reduction='sum' instead.

assert win is not None, 'Must define a window to update',
    AssertionError: Must define a window to update

in __next__:raise StopIteration  
   StopIteration

# issue #214 reference https://github.com/amdegroot/ssd.pytorch/issues/214 
try:
    images, targets = next(batch_iterator)
except StopIteration:
    batch_iterator = iter(data_loader)
images, targets = next(batch_iterator)