from data import *
from utils.augmentations import SSDAugmentation
from utils.logger import logger
from utils.helper import check_dir, resume_checkpoint, save_checkpoint
from layers.modules import MultiBoxLoss, ImageLevelLoss, HalfMultiBoxLoss
from ssd import build_ssd
# from ssd_consistency import build_ssd_con
from csd import build_ssd_con
import os
import sys
import time
import torch
from torch.autograd import Variable
from torch.utils.data import RandomSampler
import torch.nn.functional as F
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
import math
import copy
import shutil
from tensorboardX import SummaryWriter
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='VOC07', choices=['VOC07', 'VOC12', "COCO"],
type=str, help='VOC07, VOC12 or COCO')
parser.add_argument('--model', default=300, choices=[300, 512],
type=int, help='300 or 512')
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=True, type=str2bool, # None 'weights/ssd300_COCO_80000.pth'
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('--image_level_supervision', default=1, type=float,
help='whether use image level supervision')
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-5, 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('--ramp_num', default=2, type=float,
help='ramp_num')
parser.add_argument('--half_weight', default=1, type=float,
help='half_weight')
parser.add_argument('--save_folder', default='weights/',
help='Directory for saving checkpoint models')
parser.add_argument('--supervise_percent', default=0.5, type=float,
help='supervise_percent')
parser.add_argument('--supervise_num', default=0, type=int,
help='supervise_num, superior than supervise_percent')
parser.add_argument('--semi', default=True, type=str2bool,
help='semi')
parser.add_argument('--label', default=False, type=str2bool,
help='label')
parser.add_argument('--super_iter', default=1, type=int,
help='super_iter')
parser.add_argument('--add_anno', default="", type=str,
help='')
parser.add_argument('--batch_super_times', default=1, type=int,
help='batch_super_times')
args = parser.parse_args()
if torch.cuda.is_available():
if args.cuda:
torch.set_default_tensor_type('torch.cuda.FloatTensor')
if not args.cuda:
logger.info("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)
image_level_loss_weight = args.image_level_supervision
image_level_str = str(image_level_loss_weight)
if image_level_loss_weight == 1.0:
image_level_str = "1"
if image_level_loss_weight == 0.0:
image_level_str = "0"
semi_loss_weight = float(int(args.semi))
supervised = str(args.supervise_percent)
if args.supervise_num:
supervised = str(args.supervise_num)
suffix = ""
suffix ="_conf" if args.semi is True else ""
label_name = ""
if args.label:
label_name = "_label"
half_weight = args.half_weight
# checkpoint file
checkpoint_file = os.path.join(args.save_folder, "FINETUNE/data_{}_imagelevel_{}_lr_{}_ramp_{}_halfw_{}_superpercent_{}_super_time_{}_batch_{}{}{}{}" \
.format(args.dataset, image_level_str, args.lr, args.ramp_num, half_weight, supervised, args.batch_super_times, args.batch_size, suffix, label_name, args.add_anno))
check_dir(checkpoint_file)
tb_log = os.path.join(checkpoint_file, "log")
if os.path.exists(tb_log) and args.resume is False:
continue_flag = "n"
logger.info("tb_log exists, but resume is set as false, you want to remove tb_log? y/[n]")
continue_flag = input()
if continue_flag == "y":
shutil.rmtree(tb_log)
else:
logger.info("exit")
exit()
tb_writer = SummaryWriter(tb_log)
def train():
if args.dataset == 'COCO':
args.dataset_root = COCO17_ROOT
if args.model == 512:
cfg = coco512
elif args.model == 300:
cfg = coco300
else:
raise ValueError("unsupported model")
supervised_dataset = COCO17Detection(root=args.dataset_root,
supervise_percent=args.supervise_percent,
supervise_num = args.supervise_num,
only_supervise=True,
extracting_label=args.label,
finetune=True,
add_anno=args.add_anno,
transform=SSDAugmentation(cfg['min_dim'], MEANS))
if args.supervise_percent == 0:
supervised_dataset.set_only_supervise(False)
supervised_batch = args.batch_size
batch_sampler = VOCBatchSampler(RandomSampler(supervised_dataset),
batch_size=supervised_batch, drop_last=True,
super_threshold=supervised_dataset.supervise_num,
supervise_percent=supervised_dataset.supervise_percent,
only_supervise=True, batch_super_times=args.batch_super_times)
finish_flag = True
while(finish_flag):
ssd_net = build_ssd_con('train', cfg['min_dim'], cfg['num_classes'])
net = ssd_net
if args.cuda:
net = torch.nn.DataParallel(ssd_net)
cudnn.benchmark = True
init_model = os.path.join(args.save_folder, "FINETUNE", "ckpt_400000.pth")
checkpoint_model = torch.load(init_model)
model_state_dict = checkpoint_model["model_state_dict"]
optimizer_state_dict = checkpoint_model["optimizer_state_dict"]
args.start_iter = 0
net.load_state_dict(model_state_dict)
if args.cuda:
net = net.cuda()
optimizer = optim.SGD(net.parameters(), lr=args.lr, momentum=args.momentum,
weight_decay=args.weight_decay)
optimizer.load_state_dict(optimizer_state_dict)
criterion = MultiBoxLoss(cfg['num_classes'], 0.5, True, 0, True, 3, 0.5,
False, args.cuda)
image_level_criterion = ImageLevelLoss(cfg['num_classes'], args.cuda)
half_criterion = HalfMultiBoxLoss(cfg['num_classes'], 0.5, True, 0, True, 3, 0.5,
False, args.cuda)
conf_consistency_criterion = torch.nn.KLDivLoss(size_average=False, reduce=False).cuda()
adjust_learning_rate(optimizer, args.gamma, 0)
net.train()
# loss counters
loc_loss = 0
conf_loss = 0
epoch = 0
supervised_flag = 0
logger.info('Loading the dataset...')
logger.info('Training SSD on: {}'.format(args.dataset))
logger.info('Using the specified args:')
logger.info(args)
step_index = 0
total_un_iter_num = 0
#unsupervised_batch = args.batch_size - supervised_batch
#data_shuffle = 0
supervised_dataset.set_only_supervise((not args.semi))
batch_sampler.set_only_supervise((not args.semi))
supervised_data_loader = data.DataLoader(supervised_dataset,
batch_sampler=batch_sampler,
num_workers=4,
collate_fn=finetune_detection_collate,
pin_memory=True)
batch_iterator = iter(supervised_data_loader)
scale = args.batch_size / 32
for iteration in range(args.start_iter, int(200000 / scale)):
# # modify learning rate
# if iteration in cfg['lr_steps']:
# # step_index += 1
# step_index = cfg["lr_steps"].index(int(iteration * scale)) + 1
# adjust_learning_rate(optimizer, args.gamma, step_index)
# get data
try:
images, targets, semis, image_level_target, det_res = \
next(batch_iterator)
except StopIteration:
batch_iterator = iter(supervised_data_loader)
images, targets, semis, image_level_target, det_res = \
next(batch_iterator)
print("$$$$$$$$$$$$$$$$$$change data")
dd = np.array(semis)
if args.cuda:
images = Variable(images.cuda())
targets = [Variable(ann.cuda(), volatile=True) for ann in targets]
det_res = [Variable(ann.cuda(), volatile=True) for ann in det_res]
image_level_target = Variable(image_level_target.cuda())
else:
images = Variable(images)
targets = [Variable(ann, volatile=True) for ann in targets]
det_res = [Variable(ann, volatile=True) for ann in det_res]
image_level_target = Variable(image_level_target)
# forward
t0 = time.time()
# import IPython; IPython.embed(); exit()
out, conf, conf_flip, loc, loc_flip = net(images)
loss_image_level, margin, average_num = image_level_criterion(out[1], image_level_target)
sup_image_binary_index = np.zeros([len(semis),1])
half_image_index = []
full_target = []
half_target = []
half_det = []
for super_image in range(len(semis)):
if(int(semis[super_image])==1):
sup_image_binary_index[super_image] = 1
else:
sup_image_binary_index[super_image] = 0
if len(targets[super_image]) > 0:
half_target.append(targets[super_image])
half_image_index.append(super_image)
half_det.append(det_res[super_image])
# if(int(semis[len(semis)-1-super_image])==0):
# del targets[len(semis)-1-super_image]
sup_image_index = np.where(sup_image_binary_index == 1)[0]
unsup_image_index = np.where(sup_image_binary_index == 0)[0]
half_image_index = np.array(half_image_index)
for bi in sup_image_index:
full_target.append(targets[bi])
print("sup, unsup, half", (dd==1).sum(), (dd==0).sum(), len(half_target))
loc_data, conf_data, priors = out
if (len(sup_image_index) != 0):
full_loc_data = loc_data[sup_image_index,:,:]
full_conf_data = conf_data[sup_image_index,:,:]
output = (
full_loc_data,
full_conf_data,
priors
)
loss_l = Variable(torch.cuda.FloatTensor([0]))
loss_c = Variable(torch.cuda.FloatTensor([0]))
if(len(sup_image_index)!=0):
loss_l, loss_c = criterion(output, full_target)
if (len(half_det) != 0):
half_loc_data = loc_data[half_image_index,:,:]
half_conf_data = conf_data[half_image_index,:,:]
half_output = (
half_loc_data,
half_conf_data,
priors
)
half_loss_l = Variable(torch.cuda.FloatTensor([0]))
half_loss_c = Variable(torch.cuda.FloatTensor([0]))
if(len(half_target)!=0):
half_loss_l, half_loss_c = half_criterion(half_output, half_target, half_det)
# sampling = True
if(args.semi is True):
logger.info("semi")
conf_class = conf[:,:,1:].clone()
background_score = conf[:, :, 0].clone()
each_val, each_index = torch.max(conf_class, dim=2)
mask_val = each_val > background_score
mask_val = mask_val.data
mask_conf_index = mask_val.unsqueeze(2).expand_as(conf)
mask_loc_index = mask_val.unsqueeze(2).expand_as(loc)
conf_mask_sample = conf.clone()
loc_mask_sample = loc.clone()
conf_sampled = conf_mask_sample[mask_conf_index].view(-1, cfg["num_classes"])
loc_sampled = loc_mask_sample[mask_loc_index].view(-1, 4)
conf_mask_sample_flip = conf_flip.clone()
loc_mask_sample_flip = loc_flip.clone()
conf_sampled_flip = conf_mask_sample_flip[mask_conf_index].view(-1, cfg["num_classes"])
loc_sampled_flip = loc_mask_sample_flip[mask_loc_index].view(-1, 4)
# print("mask_val.sum: ", mask_val.sum())
if(args.semi and mask_val.sum()>0):
## JSD !!!!!1
conf_sampled_flip = conf_sampled_flip + 1e-7
conf_sampled = conf_sampled + 1e-7
consistency_conf_loss_a = conf_consistency_criterion(conf_sampled.log(), conf_sampled_flip.detach()).sum(-1).mean()
consistency_conf_loss_b = conf_consistency_criterion(conf_sampled_flip.log(), conf_sampled.detach()).sum(-1).mean()
consistency_conf_loss = consistency_conf_loss_a + consistency_conf_loss_b
## LOC LOSS
consistency_loc_loss_x = torch.mean(torch.pow(loc_sampled[:, 0] + loc_sampled_flip[:, 0], exponent=2))
consistency_loc_loss_y = torch.mean(torch.pow(loc_sampled[:, 1] - loc_sampled_flip[:, 1], exponent=2))
consistency_loc_loss_w = torch.mean(torch.pow(loc_sampled[:, 2] - loc_sampled_flip[:, 2], exponent=2))
consistency_loc_loss_h = torch.mean(torch.pow(loc_sampled[:, 3] - loc_sampled_flip[:, 3], exponent=2))
consistency_loc_loss = torch.div(
consistency_loc_loss_x + consistency_loc_loss_y + consistency_loc_loss_w + consistency_loc_loss_h,
4)
else:
# logger.info("not semi@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@")
consistency_conf_loss = Variable(torch.cuda.FloatTensor([0]))
consistency_loc_loss = Variable(torch.cuda.FloatTensor([0]))
consistency_loss = torch.div(consistency_conf_loss,2) + consistency_loc_loss
ramp_weight = rampweight(iteration)
consistency_loss = torch.mul(consistency_loss, ramp_weight)
# print("sup_image_index", len(sup_image_index))
if(supervised_flag >= 1):
loss = loss_l + loss_c + consistency_loss + \
image_level_loss_weight * loss_image_level + \
half_weight * (half_loss_l + half_loss_c)
else:
if(len(sup_image_index)==0):
loss = consistency_loss + image_level_loss_weight * loss_image_level + \
half_weight * (half_loss_l + half_loss_c)
# continue
else:
loss = loss_l + loss_c + consistency_loss + \
image_level_loss_weight * loss_image_level+ \
half_weight * (half_loss_l + half_loss_c)
if(loss.data>0):
optimizer.zero_grad()
loss.backward()
optimizer.step()
t1 = time.time()
if(len(sup_image_index)==0):
loss_l.data = Variable(torch.cuda.FloatTensor([0]))
loss_c.data = Variable(torch.cuda.FloatTensor([0]))
loss_image_level.data = Variable(torch.cuda.FloatTensor([0]))
logger.info("len(sup_image_index)==0")
else:
loc_loss += loss_l.data # [0]
conf_loss += loss_c.data # [0]
if iteration % 10 == 0:
logger.info('timer: %.4f sec.' % (t1 - t0))
logger.info('iter ' + repr(iteration) + ' || Loss: %.4f || consistency_loss : %.4f ||' % (loss.data, consistency_loss.data))
logger.info('loss: %.4f , loss_c: %.4f , loss_l: %.4f , loss_con: %.4f, lr : %.4f, super_len : %d\n' % (loss.data, loss_c.data, loss_l.data, consistency_loss.data,float(optimizer.param_groups[0]['lr']),len(sup_image_index)))
if iteration % 100 == 0:
tb_writer.add_scalar("Train/loss", loss.data, iteration)
tb_writer.add_scalar("Train/loss_c", loss_c.data, iteration)
tb_writer.add_scalar("Train/loss_l", loss_l.data, iteration)
tb_writer.add_scalar("Train/half_loss_c", half_loss_c.data, iteration)
tb_writer.add_scalar("Train/half_loss_l", half_loss_l.data, iteration)
tb_writer.add_scalar("Train/loss_image_level", loss_image_level.data, iteration)
tb_writer.add_scalar("Train/loss_con", consistency_loss.data, iteration)
tb_writer.add_scalar("Train/lr", float(optimizer.param_groups[0]['lr']), iteration)
tb_writer.add_scalar("Train/ramp_weight", float(ramp_weight), iteration)
tb_writer.add_scalar("Train/sup_len", float(len(sup_image_index)), iteration)
tb_writer.add_scalar("Train/margin", float(margin), iteration)
tb_writer.add_scalar("Train/average_num", float(average_num), iteration)
if(float(loss)>10000):
print("$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$")
# import IPython; IPython.embed(); exit()
break
if iteration != 0 and (iteration+1) % 1000 == 0:
# logger.info('Saving state, iter:', iteration)
# torch.save(ssd_net.state_dict(), 'weights/ssd300_COCO_' +
# repr(iteration+1) + '.pth')
save_checkpoint((iteration+1), net, optimizer, checkpoint_file)
# torch.save(ssd_net.state_dict(), args.save_folder + '' + args.dataset + '.pth')
logger.info('-------------------------------\n')
logger.info(loss.data)
logger.info('-------------------------------')
if((iteration +1) == 200000):
finish_flag = False
def rampweight(iteration):
scale = args.batch_size / 32
if args.dataset == "COCO":
scale = scale / 4 * 1.2
scale = scale * 2
ramp_up_end = 32000 / scale
ramp_down_start = 100000 / scale
ramp_end = 120000 / scale
step = 20000 / scale
if(iteration<ramp_up_end):
ramp_weight = math.exp(-5 * math.pow((1 - iteration / ramp_up_end),2))
elif(iteration>ramp_down_start):
ramp_weight = math.exp(-12.5 * math.pow((1 - (ramp_end - iteration) / step),2))
else:
ramp_weight = 1
if(iteration==0):
ramp_weight = 0
return ramp_weight / args.ramp_num
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_()
if __name__ == '__main__':
train()
