def __init__(self, viz_arch=False):
self.cfg = cfg
# Build model
print('===> Building model')
self.model, self.priorbox = create_model(cfg.MODEL)
with torch.no_grad():
self.priors = Variable(self.priorbox.forward())
# Print the model architecture and parameters
# if viz_arch is True:
# print('Model architectures:\n{}\n'.format(self.model))
# # Utilize GPUs for computation
# self.use_gpu = torch.cuda.is_available()
# self.half = False
# if self.use_gpu:
# print('Utilize GPUs for computation')
# print('Number of GPU available', torch.cuda.device_count())
# self.model.cuda()
# self.priors.cuda()
# cudnn.benchmark = True
# # self.model = torch.nn.DataParallel(self.model).module
# # Utilize half precision
# self.half = cfg.MODEL.HALF_PRECISION
# if self.half:
# self.model = self.model.half()
# self.priors = self.priors.half()
# Build preprocessor and detector
# self.preprocessor = preproc(cfg.MODEL.IMAGE_SIZE, cfg.DATASET.PIXEL_MEANS, -2)
self.detector = Detect(cfg.POST_PROCESS, self.priors)
def __init__(self):
self.cfg = cfg
self.preproc = preproc(cfg.DATASET.IMAGE_SIZE, cfg.DATASET.PIXEL_MEANS,
-2)
# Build model
print('===> Building model')
self.model, self.priorbox = create_model(cfg.MODEL)
self.priors = Variable(self.priorbox.forward(), volatile=True)
self.detector = Detect(cfg.POST_PROCESS, self.priors)
# Utilize GPUs for computation
self.use_gpu = torch.cuda.is_available()
if self.use_gpu:
print('Utilize GPUs for computation')
print('Number of GPU available', torch.cuda.device_count())
self.model.cuda()
self.priors.cuda()
cudnn.benchmark = True
# if torch.cuda.device_count() > 1:
# self.model = torch.nn.DataParallel(self.model).module
# Print the model architecture and parameters
print('Model architectures:\n{}\n'.format(self.model))
num_parameters = sum([l.nelement() for l in self.model.parameters()])
print('number of parameters: {}'.format(num_parameters))
self.max_epochs = cfg.TRAIN.MAX_EPOCHS
# Set the logger
self.output_dir = cfg.EXP_DIR
self.checkpoint = cfg.RESUME_CHECKPOINT
self.checkpoint_prefix = cfg.CHECKPOINTS_PREFIX
def test():
confg_file = '/home/hby/mycode/ssds.pytorch-1/experiments/cfgs/my_ssdlite_mobilenetv2_6.yml'
cfg_from_file(confg_file)
model, priors = create_model(cfg.MODEL)
# Utilize GPUs for computation
use_gpu = torch.cuda.is_available()
#half = False
if use_gpu:
print('Utilize GPUs for computation')
print('Number of GPU available', torch.cuda.device_count())
model.cuda()
# self.model = torch.nn.DataParallel(self.model).module
# Utilize half precision
half = cfg.MODEL.HALF_PRECISION
if half:
model = model.half()
pthfile = r'/home/hby/mycode/ssds.pytorch-1/experiments/models/ssd_mobilenet_v2_voc_6/ssd_lite_mobilenet_v2_voc_epoch_400.pth'
checkpoint = torch.load(pthfile, map_location='cuda' if use_gpu else 'cpu')
model.load_state_dict(checkpoint)
#data type nchw
dummy_input1 = torch.randn(1, 3, 300, 300).cuda().half()
# dummy_input2 = torch.randn(1, 3, 64, 64)
# dummy_input3 = torch.randn(1, 3, 64, 64)
input_names = [ "actual_input_1"]
output_names = [ "output1" ]
# torch.onnx.export(model, (dummy_input1, dummy_input2, dummy_input3), "C3AE.onnx", verbose=True, input_names=input_names, output_names=output_names)
torch.onnx.export(model, dummy_input1, "ssdmobilenet_half.onnx", verbose=True, input_names=input_names, output_names=output_names)
def __init__(self):
self.cfg = cfg
# Load data
print('===> Loading data')
self.train_loader = load_data(
cfg.DATASET, 'train') if 'train' in cfg.PHASE else None
self.eval_loader = load_data(cfg.DATASET,
'eval') if 'eval' in cfg.PHASE else None
self.test_loader = load_data(cfg.DATASET,
'test') if 'test' in cfg.PHASE else None
self.visualize_loader = load_data(
cfg.DATASET, 'visualize') if 'visualize' in cfg.PHASE else None
# Build model
print('===> Building model')
self.model, self.priorbox = create_model(cfg.MODEL)
with torch.no_grad():
self.priors = torch.Tensor(self.priorbox.forward())
self.detector = Detect(cfg.POST_PROCESS, self.priors)
# Utilize GPUs for computation
self.use_gpu = torch.cuda.is_available()
self.device = torch.device("cuda:0" if self.use_gpu else "cpu")
if self.use_gpu:
print('Utilize GPUs for computation')
print('Number of GPU available', torch.cuda.device_count())
# self.model.cuda()
# self.priors.cuda()
self.model = self.model.to(device)
self.priors = self.priors.to(device)
cudnn.benchmark = True
# if torch.cuda.device_count() > 1:
# self.model = torch.nn.DataParallel(self.model).module
# Print the model architecture and parameters
print('Model architectures:\n{}\n'.format(self.model))
# print('Parameters and size:')
# for name, param in self.model.named_parameters():
# print('{}: {}'.format(name, list(param.size())))
# print trainable scope
print('Trainable scope: {}'.format(cfg.TRAIN.TRAINABLE_SCOPE))
trainable_param = self.trainable_param(cfg.TRAIN.TRAINABLE_SCOPE)
self.optimizer = self.configure_optimizer(trainable_param,
cfg.TRAIN.OPTIMIZER)
self.exp_lr_scheduler = self.configure_lr_scheduler(
self.optimizer, cfg.TRAIN.LR_SCHEDULER)
self.max_epochs = cfg.TRAIN.MAX_EPOCHS
# metric
self.criterion = MultiBoxLoss(cfg.MATCHER, self.priors, self.use_gpu)
# Set the logger
self.writer = SummaryWriter(log_dir=cfg.LOG_DIR)
self.output_dir = cfg.EXP_DIR
self.checkpoint = cfg.RESUME_CHECKPOINT
self.checkpoint_prefix = cfg.CHECKPOINTS_PREFIX
def __init__(self, input_cfg=None, root_dir=None, viz_arch=False):
global cfg
self.cfg = cfg if input_cfg is None else input_cfg
cfg = self.cfg
# Build model
logger.info('===> Building model')
self.model, self.priorbox = create_model(cfg.MODEL)
self.priors = self.priorbox.forward()
# Print the model architecture and parameters
if viz_arch is True:
logger.info('Model architectures:\n{}\n'.format(self.model))
# Utilize GPUs for computation
self.use_gpu = torch.cuda.is_available()
#self.half = False
self.half = cfg.MODEL.HALF_PRECISION
if self.half:
self.model = BN_convert_float(self.model.half())
self.priors = self.priors.half()
if self.use_gpu:
logger.info('Utilize GPUs for computation')
logger.info('Number of GPU available: {}'.format(torch.cuda.device_count()))
self.model.cuda()
self.priors.cuda()
cudnn.benchmark = True
# self.model = torch.nn.DataParallel(self.model).module
# Utilize half precision
# Build preprocessor and detector
self.preprocessor = preproc(cfg.MODEL.IMAGE_SIZE, cfg.DATASET.PIXEL_MEANS, -2)
self.detector = Detect(cfg.POST_PROCESS, self.priors)
# Load weight:
if cfg.RESUME_CHECKPOINT == '':
AssertionError('RESUME_CHECKPOINT can not be empty')
if cfg.RESUME_CHECKPOINT == 'latest':
checkpoint_file = os.path.join(cfg.EXP_DIR, 'checkpoint_list.txt')
assert os.path.exists(checkpoint_file), \
'RESUME_CHECKPOINT set to \'latest\' but {} does not exist'.format(checkpoint_file)
with open(checkpoint_file, 'r') as f:
last_line = f.readlines()[-1]
checkpoint_file = last_line[last_line.find(':') + 2:-1]
else:
checkpoint_file = cfg.RESUME_CHECKPOINT
if root_dir is not None:
checkpoint_file = os.path.join(root_dir, checkpoint_file)
logger.info('=> loading checkpoint {:s}'.format(checkpoint_file))
checkpoint = torch.load(checkpoint_file)
self.model.load_state_dict(checkpoint)
# test only
self.model.eval()
def __init__(self):
self.cfg = cfg
# Load data
logger.info('Loading data')
self.loaders = {}
if isinstance(cfg.PHASE, (str, unicode)):
cfg.PHASE = [cfg.PHASE]
for phase in cfg.PHASE:
self.loaders[phase] = load_data(cfg.DATASET, phase)
# Build model
logger.info("Building model...")
self.model, self.priorbox = create_model(cfg.MODEL)
with torch.no_grad():
self.priors = self.priorbox.forward()
self.detector = Detect(cfg.POST_PROCESS, self.priors)
# Utilize GPUs for computation
self.use_gpu = torch.cuda.is_available()
if self.use_gpu:
logger.info('Utilize GPUs for computation')
logger.info('Number of GPU available: {}'.format(
torch.cuda.device_count()))
self.model.cuda()
self.priors.cuda()
cudnn.benchmark = True
# if torch.cuda.device_count() > 1:
# self.model = torch.nn.DataParallel(self.model).module
# Print the model architecture and parameters
logger.info('Model architectures:\n{}\n'.format(self.model))
logger.debug('Parameters and size:')
for name, param in self.model.named_parameters():
logger.debug('{}: {}'.format(name, list(param.size())))
# print trainable scope
logger.debug('Trainable scope: {}'.format(cfg.TRAIN.TRAINABLE_SCOPE))
trainable_param = self.trainable_param(cfg.TRAIN.TRAINABLE_SCOPE)
self.optimizer = self.configure_optimizer(trainable_param,
cfg.TRAIN.OPTIMIZER)
self.exp_lr_scheduler = self.configure_lr_scheduler(
self.optimizer, cfg.TRAIN.LR_SCHEDULER)
self.max_epochs = cfg.TRAIN.MAX_EPOCHS
# metric
loss_func = FocalLoss if cfg.MATCHER.USE_FOCAL_LOSS else MultiBoxLoss
self.criterion = loss_func(cfg.MATCHER, self.priors, self.use_gpu)
# Set the logger
self.writer = SummaryWriter(log_dir=cfg.LOG_DIR)
self.output_dir = cfg.EXP_DIR
self.checkpoint = cfg.RESUME_CHECKPOINT
self.checkpoints_kept = cfg.TRAIN.CHECKPOINTS_KEPT
self.checkpoint_prefix = cfg.CHECKPOINTS_PREFIX
def __init__(self, viz_arch=False):
self.cfg = cfg
# Build model
print('===> Building model')
self.model, self.priorbox = create_model(cfg.MODEL)
self.priors = Variable(self.priorbox.forward(), volatile=True)
# Print the model architecture and parameters
if viz_arch is True:
print('Model architectures:\n{}\n'.format(self.model))
# Utilize GPUs for computation
self.use_gpu = torch.cuda.is_available()
self.device = torch.device(
'gpu') if torch.cuda.is_available() else torch.device('cpu')
self.half = False
if self.use_gpu:
print('Utilize GPUs for computation')
print('Number of GPU available', torch.cuda.device_count())
self.model.cuda()
self.priors.cuda()
cudnn.benchmark = True
# self.model = torch.nn.DataParallel(self.model).module
# Utilize half precision
self.half = cfg.MODEL.HALF_PRECISION
if self.half:
self.model = self.model.half()
self.priors = self.priors.half()
# Build preprocessor and detector
self.preprocessor = preproc(cfg.MODEL.IMAGE_SIZE,
cfg.DATASET.PIXEL_MEANS, -2)
self.detector = Detect(cfg.POST_PROCESS, self.priors)
# Load weight:
if cfg.RESUME_CHECKPOINT == '':
AssertionError('RESUME_CHECKPOINT can not be empty')
print('=> loading checkpoint {:s}'.format(cfg.RESUME_CHECKPOINT))
# checkpoint = torch.load(cfg.RESUME_CHECKPOINT)
checkpoint = torch.load(cfg.RESUME_CHECKPOINT,
map_location='gpu' if self.use_gpu else 'cpu')
self.model.load_state_dict(checkpoint)
# test only
self.model.eval()
def __init__(self):
self.cfg = cfg
# Load data
print('===> Loading data')
self.train_loader = load_data(cfg.DATASET, 'train') if 'train' in cfg.PHASE else None
self.eval_loader = load_data(cfg.DATASET, 'eval') if 'eval' in cfg.PHASE else None
self.test_loader = load_data(cfg.DATASET, 'test') if 'test' in cfg.PHASE else None
self.visualize_loader = load_data(cfg.DATASET, 'visualize') if 'visualize' in cfg.PHASE else None
if self.train_loader and hasattr(self.train_loader.dataset, "num_classes"):
cfg.POST_PROCESS.NUM_CLASSES = cfg.MATCHER.NUM_CLASSES=cfg.MODEL.NUM_CLASSES=self.train_loader.dataset.num_classes
elif self.eval_loader and hasattr(self.eval_loader.dataset, "num_classes"):
cfg.POST_PROCESS.NUM_CLASSES = cfg.MATCHER.NUM_CLASSES=cfg.MODEL.NUM_CLASSES=self.eval_loader.dataset.num_classes
elif self.test_loader and hasattr(self.test_loader.dataset, "num_classes"):
cfg.POST_PROCESS.NUM_CLASSES = cfg.MATCHER.NUM_CLASSES=cfg.MODEL.NUM_CLASSES = self.test_loader.dataset.num_classes
elif self.visualize_loader and hasattr(self.visualize_loader.dataset, "num_classes"):
cfg.POST_PROCESS.NUM_CLASSES = cfg.MATCHER.NUM_CLASSES=cfg.MODEL.NUM_CLASSES = self.visualize_loader.dataset.num_classes
# Build model
print('===> Building model, num_classes is '+str(cfg.MODEL.NUM_CLASSES))
self.model, self.priorbox = create_model(cfg.MODEL,cfg.LOSS.CONF_DISTR)
self.priors = Variable(self.priorbox.forward(), volatile=True)
self.detector = Detect(cfg.POST_PROCESS, self.priors)
# Utilize GPUs for computation
self.use_gpu = torch.cuda.is_available()
if self.use_gpu:
print('Utilize GPUs for computation')
print('Number of GPU available', torch.cuda.device_count())
self.model.cuda()
self.priors.cuda()
if torch.cuda.device_count() > 1:
print('-----DataParallel-----------')
self.model = torch.nn.DataParallel(self.model)
self.model.cuda()
#self.dp_model = torch.nn.DataParallel(self.model)
#self.model = torch.nn.DataParallel(self.model).module
#self.model = self.dp_model.module
cudnn.benchmark = True
# Print the model architecture and parameters
#print('Model architectures:\n{}\n'.format(self.model))
# print('Parameters and size:')
# for name, param in self.model.named_parameters():
# print('{}: {}'.format(name, list(param.size())))
# print trainable scope
print('Trainable scope: {}'.format(cfg.TRAIN.TRAINABLE_SCOPE))
trainable_param = self.trainable_param(cfg.TRAIN.TRAINABLE_SCOPE)
self.optimizer = self.configure_optimizer(trainable_param, cfg.TRAIN.OPTIMIZER)
self.exp_lr_scheduler = self.configure_lr_scheduler(self.optimizer, cfg.TRAIN.LR_SCHEDULER)
self.max_epochs = cfg.TRAIN.MAX_EPOCHS
# metric
#self.criterion = MultiBoxLoss(cfg.MATCHER, self.priors, self.use_gpu)
self.criterion = FocalLoss(cfg.MATCHER, self.priors, self.use_gpu, cfg.LOSS)
# Set the logger
self.writer = SummaryWriter(log_dir=cfg.LOG_DIR)
self.output_dir = cfg.EXP_DIR
self.checkpoint = cfg.RESUME_CHECKPOINT
self.pretrained= cfg.PRETRAINED
self.checkpoint_prefix = cfg.CHECKPOINTS_PREFIX
def __init__(self):
self.cfg = cfg
# set up logging
if not os.path.exists(cfg.LOG_DIR):
os.mkdir(cfg.LOG_DIR)
# utils.setup_logging(os.path.join(cfg.LOG_DIR, "log.txt"))
logging.basicConfig(filename=os.path.join(cfg.LOG_DIR, "log.txt"),
level=logging.INFO)
# Load data
logging.info('===> Loading data, phase %s' % cfg.PHASE)
self.train_loader = load_data(
cfg.DATASET, 'train') if 'train' in cfg.PHASE else None
self.eval_loader = load_data(cfg.DATASET,
'eval') if 'eval' in cfg.PHASE else None
self.test_loader = load_data(cfg.DATASET,
'test') if 'test' in cfg.PHASE else None
self.visualize_loader = load_data(
cfg.DATASET, 'visualize') if 'visualize' in cfg.PHASE else None
# Build model
logging.info('===> Building model')
self.model, self.priorbox = create_model(cfg.MODEL)
self.priors = self.priorbox.forward()
self.detector = Detect(cfg.POST_PROCESS, self.priors)
# Utilize GPUs for computation
self.use_gpu = torch.cuda.is_available()
if self.use_gpu:
logging.info('Utilize GPUs for computation')
logging.info('Number of GPU available: %d' %
torch.cuda.device_count())
self.model.cuda()
self.priors.cuda()
cudnn.benchmark = True
# if torch.cuda.device_count() > 1:
# self.model = torch.nn.DataParallel(self.model).module
# Print the model architecture and parameters
# logging.info('Model architectures:\n{}\n'.format(self.model))
# logging.info('Parameters and size:')
# for name, param in self.model.named_parameters():
# logging.info('{}: {}'.format(name, list(param.size())))
# logging.info trainable scope
logging.info('Trainable scope: {}'.format(cfg.TRAIN.TRAINABLE_SCOPE))
trainable_param = self.trainable_param(cfg.TRAIN.TRAINABLE_SCOPE)
self.optimizer = self.configure_optimizer(trainable_param,
cfg.TRAIN.OPTIMIZER)
self.exp_lr_scheduler = self.configure_lr_scheduler(
self.optimizer, cfg.TRAIN.LR_SCHEDULER)
self.max_epochs = cfg.TRAIN.MAX_EPOCHS
# metric
self.criterion = MultiBoxLoss(cfg.MATCHER, self.priors, self.use_gpu)
# Set the logger
self.writer = SummaryWriter(log_dir=cfg.LOG_DIR)
self.output_dir = cfg.EXP_DIR
self.checkpoint = cfg.RESUME_CHECKPOINT
self.checkpoint_prefix = cfg.CHECKPOINTS_PREFIX
def __init__(self, ifTrain=True):
self.cfg = cfg
# Load data
print('===> Loading data')
self.ifTrain = ifTrain
if self.ifTrain:
self.train_loader = load_data(
cfg.DATASET, 'train') if 'train' in cfg.PHASE else None
#self.eval_loader = load_data(cfg.DATASET, 'eval') if 'eval' in cfg.PHASE else None
else:
test_image_dir = os.path.join('./data/', 'ship_test_v2')
# transforms = transform.Compose([transform.Lambda(lambda x: cv2.cvtColor(np.asarray(x),cv2.COLOR_RGB2BGR)),transform.Resize([300,300]), transform.ToTensor()])
# test_set = torchvision.datasets.ImageFolder(test_image_dir, transform = transforms)
# self.test_loader = torch.utils.data.DataLoader(test_set,batch_size=8,shuffle=False,num_workers=8)
self.train_loader = load_data(
cfg.DATASET, 'train') if 'train' in cfg.PHASE else None
#self.test_loader = load_data(cfg.DATASET, 'test') if 'test' in cfg.PHASE else None
self.visualize_loader = load_data(
cfg.DATASET, 'visualize') if 'visualize' in cfg.PHASE else None
# Build model
print('===> Building model')
self.model, self.priorbox = create_model(cfg.MODEL)
self.priors = Variable(self.priorbox.forward(), volatile=True)
self.detector = Detect(cfg.POST_PROCESS, self.priors)
# Utilize GPUs for computation
self.use_gpu = torch.cuda.is_available()
#self.use_gpu = False
if self.use_gpu:
print('Utilize GPUs for computation')
print('Number of GPU available', torch.cuda.device_count())
self.model.cuda()
self.priors.cuda()
cudnn.benchmark = True
if torch.cuda.device_count() > 1:
self.model = torch.nn.DataParallel(self.model).module
#device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
#device = torch.device("cpu")
#self.model = self.model.to(device)
# Print the model architecture and parameters
print('Model architectures:\n{}\n'.format(self.model))
# print('Parameters and size:')
# for name, param in self.model.named_parameters():
# print('{}: {}'.format(name, list(param.size())))
# print trainable scope
print('Trainable scope: {}'.format(cfg.TRAIN.TRAINABLE_SCOPE))
trainable_param = self.trainable_param(cfg.TRAIN.TRAINABLE_SCOPE)
# print('trainable_param ', trainable_param)
self.optimizer = self.configure_optimizer(trainable_param,
cfg.TRAIN.OPTIMIZER)
self.exp_lr_scheduler = self.configure_lr_scheduler(
self.optimizer, cfg.TRAIN.LR_SCHEDULER)
self.max_epochs = cfg.TRAIN.MAX_EPOCHS
# metric
# print('priors ', self.priors)
self.criterion = MultiBoxLoss(cfg.MATCHER, self.priors, self.use_gpu)
# Set the logger
self.writer = SummaryWriter(log_dir=cfg.LOG_DIR)
self.output_dir = cfg.EXP_DIR
self.checkpoint = cfg.RESUME_CHECKPOINT
self.checkpoint_prefix = cfg.CHECKPOINTS_PREFIX
def __init__(self):
self.cfg = cfg
self.use_gpu = torch.cuda.is_available()
if self.use_gpu:
torch.set_default_tensor_type('torch.cuda.FloatTensor')
# Load data
print('===> Loading data')
if 'train_mimic' == cfg.PHASE[0] or 'train' == cfg.PHASE[0]:
self.train_loader_1 = load_data(cfg.DATASET, 'train')
print(cfg.DATASET.DATASET, len(self.train_loader_1))
if cfg.DATASET2.DATASET in cfg.DATASET.DATASETS:
self.train_loader_2 = load_data(cfg.DATASET2, 'train')
print(cfg.DATASET2.DATASET, len(self.train_loader_1))
else:
self.train_loader_2 = None
self.test_loader = load_data(cfg.DATASET,
'test') if 'test' in cfg.PHASE1 else None
self.corr_loader = load_data(
cfg.DATASET,
'correlation') if 'correlation' in cfg.PHASE1 else None
print('===> Building model')
self.model, self.priorbox, feature_maps = create_model(cfg.MODEL)
sizes = []
boxes = []
for maps in feature_maps:
sizes.append(maps[0] * maps[1])
for box in cfg.MODEL.ASPECT_RATIOS:
boxes.append(len(box) * 2)
self.priors = Variable(self.priorbox.forward(), volatile=True)
self.detector = Detect_fast(cfg.POST_PROCESS, self.priors)
self.Discriminator = Discriminator(cfg.LOG_DIR, cfg.MODEL.NETS,
cfg.MODEL_MIMIC.NETS,
cfg.DISCTRIMINATOR, sizes, boxes)
self.Correlation = Correlation(cfg.CORRELATION, sizes, boxes)
self.Trainer = Trainer()
self.Traier_mimic = Traier_mimic(cfg.TRAIN_MIMIC, sizes, boxes,
cfg.DISCTRIMINATOR.TYPE)
self.Tester = Tester(cfg.POST_PROCESS)
if 'train_mimic' == cfg.PHASE[0] or 'correlation' == cfg.PHASE[0]:
self.model_mimic = create_model(cfg.MODEL_MIMIC)
self.model_mimic.load_state_dict(
torch.load(cfg.MODEL_MIMIC.WEIGHTS))
self.DNet = self.Discriminator.create_discriminator()
# Utilize GPUs for computation
if self.use_gpu:
print('Utilize GPUs for computation')
print('Number of GPU available', torch.cuda.device_count())
cudnn.benchmark = True
# self.model_mimic = torch.nn.DataParallel(self.model_mimic)
# for i in range(len(self.DNet)):
# self.DNet[i] = torch.nn.DataParallel(self.DNet[i])
self.model.cuda()
self.priors.cuda()
if 'train_mimic' == cfg.PHASE[0] or 'correlation' == cfg.PHASE[0]:
self.model_mimic.cuda()
for i in range(len(self.DNet)):
self.DNet[i] = self.DNet[i].cuda()
# if 'train_mimic' == cfg.PHASE[0]:
# print('Model mimim architectures:\n{}\n'.format(self.model_mimic))
# for i in range(len(self.DNet)):
# print('Hello')
# print(self.DNet[i])
# print('Parameters and size:')
# for name, param in self.model.named_parameters():
# print('{}: {}'.format(name, list(param.size())))
# print trainable scope
# print('Trainable scope: {}'.format(cfg.TRAIN.TRAINABLE_SCOPE))
if 'train_mimic' == cfg.PHASE[0] or 'correlation' == cfg.PHASE[0]:
self.optimizer = self.configure_optimizer(self.model,
cfg.TRAIN.OPTIMIZER)
self.DNet_optim = []
for i in range(len(self.DNet)):
self.DNet_optim.append(
self.configure_optimizer(self.DNet[i],
cfg.DISCTRIMINATOR.OPTIMIZER))
self.optimizer_GENERATOR = self.configure_optimizer(
self.model, cfg.TRAIN_MIMIC.OPTIMIZER)
self.exp_lr_scheduler_g = self.configure_lr_scheduler(
self.optimizer_GENERATOR, cfg.TRAIN.LR_SCHEDULER)
else:
self.optimizer = self.configure_optimizer(self.model,
cfg.TRAIN.OPTIMIZER)
self.phase = cfg.PHASE
self.exp_lr_scheduler = self.configure_lr_scheduler(
self.optimizer, cfg.TRAIN.LR_SCHEDULER)
self.max_epochs = cfg.TRAIN.MAX_EPOCHS
# metric
self.criterion = MultiBoxLoss(cfg.MATCHER, self.priors, self.use_gpu)
self.criterion_GaN = nn.BCELoss()
self.pos = POSdata(cfg.MATCHER, self.priors, self.use_gpu)
# Set the logger
self.writer = SummaryWriter(log_dir=cfg.LOG_DIR)
if not os.path.exists(cfg.LOG_DIR):
os.mkdir(cfg.LOG_DIR)
shutil.copyfile('./lib/utils/config_parse.py',
cfg.LOG_DIR + 'hiperparameters.py')
a = os.listdir(cfg.LOG_DIR)
for i in range(1, 100):
if not 'Correlation_' + str(i) + '.txt' in a:
self.logger = cfg.LOG_DIR + 'Correlation_' + str(i) + '.txt'
self.loglosses = cfg.LOG_DIR + 'Correlation_loss_' + str(
i) + '.txt'
break
f = open(self.logger, 'w')
f.close()
f = open(self.loglosses, 'w')
f.close()
self.output_dir = cfg.LOG_DIR
self.checkpoint = cfg.RESUME_CHECKPOINT
self.checkpoint_prefix = cfg.CHECKPOINTS_PREFIX
self.model.loc.apply(self.weights_init)
self.model.conf.apply(self.weights_init)
self.model.extras.apply(self.weights_init)
if 'train' == cfg.PHASE[0]:
if torch.cuda.device_count() > 1:
self.model = torch.nn.DataParallel(self.model)
def run_detection(data_dir, coco_gt, im_ids):
model, priorbox = create_model(cfg.MODEL)
priors = Variable(priorbox.forward(), volatile=True)
detector = Detect(cfg.POST_PROCESS, priors)
# Utilize GPUs for computation
model.cuda()
priors.cuda()
cudnn.benchmark = True
preprocess = preproc(cfg.DATASET.IMAGE_SIZE, cfg.DATASET.PIXEL_MEANS, -2)
resume_checkpoint(model, args.weight)
num_classes = 2
results = []
time_all = []
time_per_step = {
"nms_time": [],
"cpu_tims": [],
"scores_time": [],
"box_time": [],
"gpunms_time": [],
"base_time": [],
"extra_time": [],
"head_time": []
}
for i, index in enumerate(im_ids):
# load img
print('evaluating image {}/{}'.format(i, len(im_ids)))
im_data = coco_gt.loadImgs(ids=index)[0]
img = cv2.imread(os.path.join(data_dir, 'frames',
im_data['file_name']))
scale = [img.shape[1], img.shape[0], img.shape[1], img.shape[0]]
img_shape = img.shape
images = Variable(preprocess(img)[0].unsqueeze(0).cuda(),
volatile=True)
img_dict = {'version': 0, 'time': 0., 'camera_id': 0, 'image': img}
# run detection model
torch.cuda.synchronize()
time_all_start = time.perf_counter()
# forward
out = model(images, phase='eval')
# detect
detections = detector.forward(out)
torch.cuda.synchronize()
time_all_end = time.perf_counter()
time_all.append(1000 * (time_all_end - time_all_start))
scores = []
cls_boxes = []
for det in detections[0][1]:
if det[0] > 0:
d = det.cpu().numpy()
score, box = d[0], d[1:]
box *= scale
scores.append(score)
cls_boxes.append(box)
#print(score)
#print(box)
output_blobs = {}
output_blobs['scores'] = scores
output_blobs['cls_boxes'] = cls_boxes
print(np.array(cls_boxes).shape)
output_dict = postprocess(output_blobs, 1., img_shape)
if len(output_dict['people']) == 0:
continue
# save result
entry_index = 0
for person in output_dict['people']:
entry_result = {
"image_id": index,
"category_id": 1,
"bbox": person['bbox_ltwh'].tolist(),
"score": person['score']
}
results.append(entry_result)
# save results as json file
with open(json_dt, 'w') as f:
json.dump(results, f)
print('detection results saved in {}'.format(json_dt))
print('average running time: {}ms'.format(sum(time_all) / len(time_all)))
def train_epoch(self, model, data_loader, optimizer, criterion, writer,
epoch, use_gpu):
model.train()
epoch_size = len(data_loader)
batch_iterator = iter(data_loader)
loc_loss = 0
conf_loss = 0
l2_loss = 0
l2_loss_weight = 0.7
_t = Timer()
#prepare for teacher model
teacher_model, teacher_priorbox = create_model(cfg.TEACHER_MODEL)
#teacher_priors = Variable(teacher_priorbox.forward(), volatile=True)
#teacher_detector = Detect(cfg.POST_PROCESS, teacher_priors)
# Utilize GPUs for computation
teacher_model.cuda()
#teacher_priors.cuda()
cudnn.benchmark = True
#teacher_preprocess = preproc(cfg.DATASET.IMAGE_SIZE, cfg.DATASET.PIXEL_MEANS, -2)
weight = '/home/zhihuai/yolo/experiments/models/darknet_53_yolo_v3_aifi_three_60ep/yolo_v3_darknet_53_coco_epoch_60.pth'
self.resume_checkpoint_teacher(teacher_model, weight)
for iteration in iter(range((epoch_size))):
images, targets = next(batch_iterator)
if use_gpu:
images = Variable(images.cuda())
targets = [
Variable(anno.cuda(), volatile=True) for anno in targets
]
else:
images = Variable(images)
targets = [Variable(anno, volatile=True) for anno in targets]
_t.tic()
# forward
out, student_out = model(images, phase='distill_student')
teacher_out = teacher_model(images, phase='distill_teacher')
teacher_feature_1 = teacher_out[0]
student_feature_1 = student_out[0]
teacher_feature_1.detach_()
student_feature_1.detach_()
loss_l2_1 = self.l2_loss(student_feature_1, teacher_feature_1)
teacher_feature_2 = teacher_out[1]
student_feature_2 = student_out[1]
teacher_feature_2.detach_()
student_feature_2.detach_()
loss_l2_2 = self.l2_loss(student_feature_2, teacher_feature_2)
# backprop
optimizer.zero_grad()
loss_l, loss_c = criterion(out, targets)
# some bugs in coco train2017. maybe the annonation bug.
if loss_l.data[0] == float("Inf"):
continue
loss = loss_l + loss_c + l2_loss_weight * loss_l2_1 + (
1 - l2_loss_weight) * loss_l2_2
loss.backward()
optimizer.step()
time = _t.toc()
loc_loss += loss_l.data[0]
conf_loss += loss_c.data[0]
l2_loss += l2_loss_weight * loss_l2_1.data[0] + (
1 - l2_loss_weight) * loss_l2_2.data[0]
# log per iter
log = '\r==>Train: || {iters:d}/{epoch_size:d} in {time:.3f}s [{prograss}] || loc_loss: {loc_loss:.4f} cls_loss: {cls_loss:.4f} l2_loss: {l2_loss:.4f}\r'.format(
prograss='#' * int(round(10 * iteration / epoch_size)) +
'-' * int(round(10 * (1 - iteration / epoch_size))),
iters=iteration,
epoch_size=epoch_size,
time=time,
loc_loss=loss_l.data[0],
cls_loss=loss_c.data[0],
l2_loss=l2_loss_weight * loss_l2_1.data[0] +
(1 - l2_loss_weight) * loss_l2_2.data[0])
sys.stdout.write(log)
sys.stdout.flush()
# log per epoch
sys.stdout.write('\r')
sys.stdout.flush()
lr = optimizer.param_groups[0]['lr']
log = '\r==>Train: || Total_time: {time:.3f}s || loc_loss: {loc_loss:.4f} conf_loss: {conf_loss:.4f} l2_loss: {l2_loss:.4f} || lr: {lr:.6f}\n'.format(
lr=lr,
time=_t.total_time,
loc_loss=loc_loss / epoch_size,
conf_loss=conf_loss / epoch_size,
l2_loss=l2_loss / epoch_size)
sys.stdout.write(log)
sys.stdout.flush()
# log for tensorboard
writer.add_scalar('Train/loc_loss', loc_loss / epoch_size, epoch)
writer.add_scalar('Train/conf_loss', conf_loss / epoch_size, epoch)
writer.add_scalar('Train/lr', lr, epoch)
def __init__(self, phase="train"):
self.cfg = cfg
creat_log(self.cfg, phase=phase)
for k, v in cfg.items():
print(k, ": ", v)
log_str = '\rEpoch {k}: {v}'.format(k=k, v=v)
logging.info(log_str)
# Load data
print('===> Loading data')
logging.info('===> Loading data')
self.train_loader = load_data(
cfg.DATASET, 'train') if 'train' in cfg.PHASE else None
self.eval_loader = load_data(cfg.DATASET,
'eval') if 'eval' in cfg.PHASE else None
self.test_loader = load_data(cfg.DATASET,
'test') if 'test' in cfg.PHASE else None
self.visualize_loader = load_data(
cfg.DATASET, 'visualize') if 'visualize' in cfg.PHASE else None
# Build model
print('===> Building model')
logging.info('===> Building model')
self.model, self.priorbox = create_model(cfg.MODEL)
with torch.no_grad():
self.priors = Variable(self.priorbox.forward())
self.detector = Detect(cfg.POST_PROCESS, self.priors)
os.makedirs(self.cfg['EXP_DIR'], exist_ok=True)
# Utilize GPUs for computation
self.use_gpu = torch.cuda.is_available()
print('Model architectures:\n{}\n'.format(self.model))
logging.info('Model architectures:\n{}\n'.format(self.model))
from lib.utils.torchsummary import summary
summary_text = summary(
self.model.cuda(),
(3, cfg.MODEL.IMAGE_SIZE[0], cfg.MODEL.IMAGE_SIZE[1]))
logging.info('\n'.join(summary_text))
# num_params = 0
# for name, param in self.model.named_parameters():
# num_params += param.numel()
# # print(name, param.size(), param.numel())
# print("%40s %20s %20s" % (name, num_params, param.numel()))
# print(num_params/1e4)
# df = torch_summarize_df(input_size=(3, 512, 512), model=self.model)
# # df['name'], list(df['class_name']), df['input_shape'], df["output_shape"], list(df['nb_params'])
# print(df)
# for name, param in self.model.named_parameters():
# print(name, param.size())
# from thop import profile
#
# flops, params = profile(self.model, input_size=(1, 3, 512, 128))
# count = 0
# for p in self.model.parameters():
# count += p.data.nelement()
# self.multi_gpu = True
self.multi_gpu = False
if self.use_gpu:
print('Utilize GPUs for computation')
logging.info('Utilize GPUs for computation')
# print('Number of GPU available', torch.cuda.device_count())
self.model.cuda()
self.priors.cuda()
cudnn.benchmark = True
# os.environ['CUDA_VISIBLE_DEVICES'] = "4,5,6,7" # "0,1,2,3,4,5,6,7"
if torch.cuda.device_count() > 1 and self.multi_gpu:
self.model = torch.nn.DataParallel(self.model.cuda())
cudnn.benchmark = True
# self.model = torch.nn.DataParallel(self.model).module
# Print the model architecture and parameters
# print('Parameters and size:')
# for name, param in self.model.named_parameters():
# print('{}: {}'.format(name, list(param.size())))
# print trainable scope
print('Trainable scope: {}'.format(cfg.TRAIN.TRAINABLE_SCOPE))
logging.info('Trainable scope: {}'.format(cfg.TRAIN.TRAINABLE_SCOPE))
trainable_param = self.trainable_param(cfg.TRAIN.TRAINABLE_SCOPE)
self.optimizer = self.configure_optimizer(trainable_param,
cfg.TRAIN.OPTIMIZER)
self.exp_lr_scheduler = self.configure_lr_scheduler(
self.optimizer, cfg.TRAIN.LR_SCHEDULER)
self.max_epochs = cfg.TRAIN.MAX_EPOCHS
# metric
self.criterion = MultiBoxLoss(cfg.MATCHER, self.priors, self.use_gpu)
# Set the logger
self.writer = SummaryWriter(logdir=cfg.LOG_DIR)
self.output_dir = cfg.EXP_DIR
self.checkpoint = cfg.RESUME_CHECKPOINT
self.checkpoint_prefix = cfg.CHECKPOINTS_PREFIX
