def run_train(args): assert torch.cuda.is_available(), 'Error: CUDA not found!' best_loss = float('inf') # best test loss start_epoch = 0 # start from epoch 0 or last epoch # Data print('==> Preparing data..') trainloader = get_train_loader(img_dir=settings.IMG_DIR, batch_size=batch_size) #trainloader = get_small_train_loader() print(trainloader.num) #testloader = get_train_loader(img_dir=settings.IMG_DIR) # Model net = RetinaNet() #net.load_state_dict(torch.load('./model/net.pth')) net.load_state_dict(torch.load('./ckps/best_0.pth')) net = torch.nn.DataParallel(net, device_ids=range(torch.cuda.device_count())) net.cuda() criterion = FocalLoss() #optimizer = optim.SGD(net.parameters(), lr=args.lr, momentum=0.9, weight_decay=1e-4) optimizer = optim.Adam(net.parameters(), lr=args.lr) iter_save = 200 bgtime = time.time() # Training for epoch in range(start_epoch, start_epoch + 100): print('\nEpoch: %d' % epoch) net.train() #net.module.freeze_bn() train_loss = 0 for batch_idx, (inputs, loc_targets, cls_targets) in enumerate(trainloader): inputs = Variable(inputs.cuda()) loc_targets = Variable(loc_targets.cuda()) cls_targets = Variable(cls_targets.cuda()) optimizer.zero_grad() loc_preds, cls_preds = net(inputs) loss = criterion(loc_preds, loc_targets, cls_preds, cls_targets) loss.backward() optimizer.step() #train_loss += loss.data[0] sample_num = (batch_idx + 1) * batch_size avg_loss = running_loss(loss.data[0]) print( 'Epoch: {}, num: {}/{} train_loss: {:.3f} | run_loss: {:.3f} min: {:.1f}' .format(epoch, sample_num, trainloader.num, loss.data[0], avg_loss, (time.time() - bgtime) / 60), end='\r') if batch_idx % iter_save == 0: torch.save( net.module.state_dict(), './ckps/best_{}.pth'.format(batch_idx // iter_save % 5)) log.info('batch: {}, loss: {:.4f}'.format(batch_idx, avg_loss))
def main(): print('==> chooseing data..') transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)) ]) # Model net = RetinaNet() criterion = FocalLoss() # optimizer = optim.Adam(net.parameters(), lr=0.001, betas=(0.9, 0.999), eps=1e-08, weight_decay=1e-4) optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9, weight_decay=1e-4) load_model_epoch = args.load_model_epoch checkpoint = torch.load( './checkpoint/{}_ckpt.pth'.format(load_model_epoch)) # max_epoch net = torch.nn.DataParallel(net, device_ids=range(torch.cuda.device_count())) net.cuda() net.load_state_dict(checkpoint['net']) optimizer.load_state_dict(checkpoint['optimizer']) start_epoch = checkpoint['epoch'] test(start_epoch, transform, net, criterion, optimizer)
def main(): # assert torch.cuda.is_available(), 'Error: CUDA not found!' best_loss = float('inf') # best test loss start_epoch = 0 # start from epoch 0 or last epoch save_model_path = args.model # Data print('==> Preparing data..') transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)) ]) # Model net = RetinaNet() net.load_state_dict(torch.load('./model/net.pth')) net = torch.nn.DataParallel(net, device_ids=range(torch.cuda.device_count())) net.cuda() criterion = FocalLoss() # optimizer = optim.Adam(net.parameters(), lr=args.lr, betas=(0.9, 0.999), eps=1e-08, weight_decay=1e-4) optimizer = optim.SGD(net.parameters(), lr=args.lr, momentum=0.9, weight_decay=1e-4) for epoch in range(start_epoch, start_epoch + args.train_epoch): train(epoch + 1, transform, net, optimizer, criterion) save_model(epoch + 1, save_model_path, net, optimizer)
def train(total_epochs=1, interval=100, resume=False, ckpt_path = ''): print("Loading training dataset...") train_dset = OpenImagesDataset(root='./data/train', list_file ='./data/tmp/train_images_bbox.csv', transform=transform, train=True, input_size=600) train_loader = data.DataLoader(train_dset, batch_size=4, shuffle=True, num_workers=4, collate_fn=train_dset.collate_fn) print("Loading completed.") #val_dset = OpenImagesDataset(root='./data/train', # list_file='./data/tmp/train_images_bbox.csv', train=False, transform=transform, input_size=600) #val_loader = torch.utils.data.DataLoader(val_dset, batch_size=1, shuffle=False, num_workers=4, collate_fn=val_dset.collate_fn) net = RetinaNet() net.load_state_dict(torch.load('./model/net.pth')) criterion = FocalLoss() net.cuda() criterion.cuda() optimizer = optim.SGD(net.parameters(), lr=1e-3, momentum=0.9, weight_decay=1e-4) best_val_loss = 1000 start_epoch=0 if resume: if os.path.isfile(ckpt_path): print(f'Loading from the checkpoint {ckpt_path}') checkpoint = torch.load(ckpt_path) start_epoch = checkpoint['epoch'] best_val_loss = checkpoint['best_val_loss'] net.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) print(f'Loaded checkpoint {ckpt_path}, epoch : {start_epoch}') else: print(f'No check point found at the path {ckpt_path}') for epoch in range(start_epoch, total_epochs): train_one_epoch(train_loader, net, criterion, optimizer, epoch, interval) val_loss = 0 #val_loss = validate(val_loader, net, criterion, interval) if val_loss < best_val_loss: best_val_loss = val_loss save_checkpoint({ 'epoch': epoch+1, 'state_dict': net.state_dict(), 'best_val_loss': best_val_loss, 'optimizer' : optimizer.state_dict() }, is_best=True)
def train_obj(): # Model net = RetinaNet() net = torch.nn.DataParallel(net, device_ids=range( torch.cuda.device_count())) net.cuda() criterion = FocalLoss() optimizer = optim.SGD(net.parameters(), lr=args.lr, momentum=0.9, weight_decay=1e-4) scheduler_obj = torch.optim.lr_scheduler.LambdaLR( optimizer, lr_lambda=lambda x: (1 - x / (len(trainloader) * epochs))**0.9) obj_trainer = ObjDetTrainer(net, criterion, optimizer, scheduler_obj, trainloader, valloader, device) obj_trainer.train(epochs, True)
if fix == 'head': for param in net.module.fpn.conv1.parameters(): param.requires_grad = False for param in net.module.fpn.bn1.parameters(): param.requires_grad = False for param in net.module.fpn.layer1.parameters(): param.requires_grad = False for param in net.module.fpn.layer2.parameters(): param.requires_grad = False for param in net.module.fpn.layer3.parameters(): param.requires_grad = False for param in net.module.fpn.layer4.parameters(): param.requires_grad = False else: for param in net.parameters(): param.requires_grad = True from itertools import ifilter op_parameters = ifilter(lambda p: p.requires_grad, net.parameters()) ''' count=0 for param in net.parameters(): count+=1 print(param.requires_grad) print(count) ''' criterion = FocalLoss(num_classes=15) optimizer = optim.Adam(op_parameters, lr=lr, betas=(0.9, 0.99)) #later add scheduler into optimizer from torch.optim import lr_scheduler
backbone = resnet18_features elif args.backbone == 'resnet34': backbone = resnet34_features elif args.backbone == 'resnet50': backbone = resnet50_features elif args.backbone == 'resnet101': backbone = resnet101_features elif args.backbone == 'resnet152': backbone = resnet152_features else: raise ValueError('Invalid model backbone') # RetinaNet model model = RetinaNet(80, backbone, pretrained=True) optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9, weight_decay=0.0001) # TODO: enable more flexible lr profile scheduler = lr_scheduler.MultiStepLR( optimizer, # Milestones are set assuming batch size is 16: # 60000 / batch_size = 3750 # 80000 / batch_size = 5000 milestones=[3750, 5000], gamma=0.1) focal_loss = FocalLoss(80) smoothed_l1_loss = SmoothedL1Loss(3.)
checkpoint = torch.load('/media/Darius/shayeree/training/checkpoint_list_wise/albertsons/Retina50ProdB1/ckpt_0005_0.1237.pth') net.load_state_dict(checkpoint['net']) best_loss = checkpoint['loss'] start_epoch = checkpoint['epoch'] else: params_dict = torch.load('./model/{:s}.pth'.format(args.net)) net_dict = net.fpn.state_dict() params_dict = {k: v for k, v in params_dict.items() if k in net_dict} net_dict.update(params_dict) net.fpn.load_state_dict(net_dict) #net,optimizer=amp.initialize(net,optimizer,opt_level='01',loss_scale="dynamic") #net = torch.nn.DataParallel(net, device_ids=range(torch.cuda.device_count())) net.cuda() criterion = FocalLoss(num_classes=1) optimizer = optim.SGD(net.parameters(), lr=args.lr, momentum=0.9, weight_decay=1e-4) amp_handle = amp.init(enabled=True, verbose=True) optimizer = amp_handle.wrap_optimizer(optimizer) #net,optimizer=amp.initialize(net.cuda(),optimizer,opt_level='O1',loss_scale="dynamic") #net = torch.nn.DataParallel(net, device_ids=range(torch.cuda.device_count())) #net=apex.parallel.DistributedDataParallel(net, device_ids=range(torch.cuda.device_count())) #scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=5,gamma=0.1) #scheduler = optim.lr_scheduler.MultiStepLR(optimizer, milestones=[1], gamma=0.1) # Training def train(epoch): # print('\nEpoch: %d' % epoch) net.train()
# Model net = RetinaNet() net.load_state_dict(torch.load('./model/net.pth')) args.resume = True if args.resume: print('==> Resuming from checkpoint..') checkpoint = torch.load('./checkpoint/ckpt.pth') net.load_state_dict(checkpoint['net']) # best_loss = checkpoint['loss'] start_epoch = checkpoint['epoch'] net = torch.nn.DataParallel(net, device_ids=range(torch.cuda.device_count())) net.cuda() criterion = FocalLoss() optimizer = optim.SGD(net.parameters(), lr=args.lr, momentum=0.9, weight_decay=1e-4) # Training def train(epoch): print('\nEpoch: %d' % epoch) net.train() net.module.freeze_bn() train_loss = 0 for batch_idx, (inputs, loc_targets, cls_targets, _) in enumerate(trainloader): inputs = Variable(inputs.cuda()) loc_targets = Variable(loc_targets.cuda())
transforms.Normalize(cfg.mean, cfg.std) ]) trainset = VocLikeDataset(image_dir=cfg.image_dir, annotation_dir=cfg.annotation_dir, imageset_fn=cfg.train_imageset_fn, image_ext=cfg.image_ext, classes=cfg.classes, encoder=DataEncoder(), transform=train_transform) trainloader = torch.utils.data.DataLoader(trainset, batch_size=cfg.batch_size, shuffle=True, num_workers=cfg.num_workers, collate_fn=trainset.collate_fn) print('Building model...') net = RetinaNet(backbone=cfg.backbone, num_classes=len(cfg.classes)) net = torch.nn.DataParallel(net, device_ids=range(torch.cuda.device_count())) net.cuda() criterion = FocalLoss(len(cfg.classes)) optimizer = optim.SGD(net.parameters(), lr=lr, momentum=cfg.momentum, weight_decay=cfg.weight_decay) # special_layers=nn.ModuleList(net.subnet_classes) # special_layers_params=list(map(id,special_layers.parameters())) # base_params=filter(lambda p:id(p) not in special_layers_params ,net.parameters()) # optimizer=torch.optim.SGD([{'params':base_params},{'params':special_layers_params.parameters(),'lr':0.001}],lr=lr, momentum=cfg.momentum, weight_decay=cfg.weight_decay) def save_checkpoint(loss, net, n): global best_loss loss /= n if loss < best_loss: print('Saving..') state = { 'net': net.state_dict(), 'loss': loss,
net = RetinaNet(backbone=cfg.backbone, num_classes=len(cfg.classes)) net = torch.nn.DataParallel(net, device_ids=range(torch.cuda.device_count())) net.cuda() cudnn.benchmark = True if args.resume: print('Resuming from checkpoint..') checkpoint = torch.load(os.path.join('ckpts', args.exp, '30_ckpt.pth')) net.load_state_dict(checkpoint['net']) start_epoch = checkpoint['epoch'] lr = cfg.lr criterion = FocalLoss(len(cfg.classes)) optimizer = optim.Adam(net.parameters(), lr=cfg.lr, weight_decay=cfg.weight_decay) writer = SummaryWriter(log_dir='logs') # input_data = torch.rand(cfg.batch_size, 3, cfg.width, cfg.height) # writer.add_graph(net, input_data) def train(epoch): trainloader = torch.utils.data.DataLoader(trainset, batch_size=cfg.batch_size, shuffle=True, num_workers=cfg.num_workers, collate_fn=trainset.collate_fn) print('\nTrain Epoch: %d' % epoch)
# Model net = RetinaNet() net.load_state_dict(torch.load('./model/net.pth')) if args.resume: print('==> Resuming from checkpoint..') checkpoint = torch.load('./checkpoint/ckpt.pth') net.load_state_dict(checkpoint['net']) best_loss = checkpoint['loss'] start_epoch = checkpoint['epoch'] net = torch.nn.DataParallel(net, device_ids=range(torch.cuda.device_count())) net.cuda() criterion = FocalLoss() optimizer = optim.Adam(net.parameters(), lr=args.lr, weight_decay=1e-4) print('Number of classes', criterion.num_classes) # Training def train(epoch): print('\nEpoch: %d' % epoch) net.train() train_loss = 0 for batch_idx, (inputs, loc_targets, cls_targets) in enumerate(trainloader): inputs = Variable(inputs.cuda()) loc_targets = Variable(loc_targets.cuda()) cls_targets = Variable(cls_targets.cuda())
return for c in childrens: foo(c) # resnet = models.alexnet() # resnet = models.resnet152() model = model foo(model) input = Variable(torch.rand(3, size, size).unsqueeze(0), requires_grad=True).cuda() out = model(input) total_flops = (sum(list_conv) + sum(list_linear) + sum(list_bn) + sum(list_relu) + sum(list_pooling)) print(' + Number of FLOPs: %.2fG' % (total_flops / 1e9)) if __name__ == "__main__": device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # Set up model model = RetinaNet().to(device) # summary(model, input_size=(3,opt.img_size,opt.img_size)) total = sum([param.nelement() for param in model.parameters()]) print(' + Number of params: %.2fM' % (total / 1e6)) print_model_parm_flops(model, opt.img_size)
def train(): args = parse_args() assert torch.cuda.is_available(), 'Error: CUDA not found!' assert args.focal_loss, "OHEM + ce_loss is not working... :(" if not os.path.exists(args.save_folder): os.mkdir(args.save_folder) if not os.path.exists(args.logdir): os.mkdir(args.logdir) ########################################################################### # Data ########################################################################### print('==> Preparing data..') trainset = ListDataset(root='/mnt/9C5E1A4D5E1A2116/datasets/', dataset=args.dataset, train=True, transform=Augmentation_traininig, input_size=args.input_size, multi_scale=args.multi_scale) trainloader = torch.utils.data.DataLoader(trainset, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers, collate_fn=trainset.collate_fn) ########################################################################### # Training Detail option\ stepvalues = (10000, 20000, 30000, 40000, 50000) if args.dataset in ["SynthText"] \ else (2000, 4000, 6000, 8000, 10000) best_loss = float('inf') # best test loss start_epoch = 0 # start from epoch 0 or last epoch iteration = 0 cur_lr = args.lr mean = (0.485, 0.456, 0.406) var = (0.229, 0.224, 0.225) step_index = 0 pEval = None ########################################################################### # Model ########################################################################### # set model (focal_loss vs OHEM_CE loss) if args.focal_loss: imagenet_pretrain = 'weights/retinanet_se50.pth' criterion = FocalLoss() num_classes = 1 else: imagenet_pretrain = 'weights/retinanet_se50_OHEM.pth' criterion = OHEM_loss() num_classes = 2 net = RetinaNet(num_classes) # Restore model weights net.load_state_dict(torch.load(imagenet_pretrain)) if args.resume: print('==> Resuming from checkpoint..', args.resume) checkpoint = torch.load(args.resume) net.load_state_dict(checkpoint['net']) #start_epoch = checkpoint['epoch'] #iteration = checkpoint['iteration'] #cur_lr = checkpoint['lr'] #step_index = checkpoint['step_index'] # optimizer.load_state_dict(state["optimizer"]) print("multi_scale : ", args.multi_scale) print("input_size : ", args.input_size) print("stepvalues : ", stepvalues) print("start_epoch : ", start_epoch) print("iteration : ", iteration) print("cur_lr : ", cur_lr) print("step_index : ", step_index) print("num_gpus : ", torch.cuda.device_count()) # Data parellelism for multi-gpu training net = torch.nn.DataParallel(net, device_ids=range(torch.cuda.device_count())) net.cuda() # Put model in training mode and freeze batch norm. net.train() net.module.freeze_bn() # you must freeze batchnorm ########################################################################### # Optimizer ########################################################################### optimizer = optim.SGD(net.parameters(), lr=cur_lr, momentum=0.9, weight_decay=1e-4) #optimizer = optim.Adam(net.parameters(), lr=cur_lr) ########################################################################### # Utils ########################################################################### encoder = DataEncoder() writer = SummaryWriter(log_dir=args.logdir) ########################################################################### # Training loop ########################################################################### t0 = time.time() for epoch in range(start_epoch, 10000): if iteration > args.max_iter: break for inputs, loc_targets, cls_targets in trainloader: inputs = Variable(inputs.cuda()) loc_targets = Variable(loc_targets.cuda()) cls_targets = Variable(cls_targets.cuda()) optimizer.zero_grad() loc_preds, cls_preds = net(inputs) loc_loss, cls_loss = criterion(loc_preds, loc_targets, cls_preds, cls_targets) loss = loc_loss + cls_loss loss.backward() optimizer.step() if iteration % 20 == 0: t1 = time.time() print( 'iter ' + repr(iteration) + ' (epoch ' + repr(epoch) + ') || loss: %.4f || l loc_loss: %.4f || l cls_loss: %.4f (Time : %.1f)' % (loss.sum().item(), loc_loss.sum().item(), cls_loss.sum().item(), (t1 - t0))) # t0 = time.time() writer.add_scalar('loc_loss', loc_loss.sum().item(), iteration) writer.add_scalar('cls_loss', cls_loss.sum().item(), iteration) writer.add_scalar('loss', loss.sum().item(), iteration) # show inference image in tensorboard infer_img = np.transpose(inputs[0].cpu().numpy(), (1, 2, 0)) infer_img *= var infer_img += mean infer_img *= 255. infer_img = np.clip(infer_img, 0, 255) infer_img = infer_img.astype(np.uint8) h, w, _ = infer_img.shape boxes, labels, scores = encoder.decode(loc_preds[0], cls_preds[0], (w, h)) boxes = boxes.reshape(-1, 4, 2).astype(np.int32) if boxes.shape[0] != 0: # infer_img = infer_img/np.float32(255) # print(boxes) # print( # f"infer_img prior to cv2.polylines - dtype: {infer_img.dtype}, shape: {infer_img.shape}, min: {infer_img.min()}, max: {infer_img.max()}") # print( # f"boxes prior to cv2.polylines - dtype: {boxes.dtype}, shape: {boxes.shape}, min: {boxes.min()}, max: {boxes.max()}") infer_img = cv2.polylines(infer_img.copy(), boxes, True, (0, 255, 0), 4) # print( # f"infer_img - dtype: {infer_img.dtype}, shape: {infer_img.shape}, min: {infer_img.min()}, max: {infer_img.max()}") writer.add_image('image', infer_img, iteration, dataformats="HWC") writer.add_scalar('input_size', h, iteration) writer.add_scalar('learning_rate', cur_lr, iteration) t0 = time.time() if iteration % args.save_interval == 0 and iteration > 0: print('Saving state, iter : ', iteration) state = { 'net': net.module.state_dict(), "optimizer": optimizer.state_dict(), 'iteration': iteration, 'epoch': epoch, 'lr': cur_lr, 'step_index': step_index } model_file = args.save_folder + \ 'ckpt_' + repr(iteration) + '.pth' torch.save(state, model_file) if iteration in stepvalues: step_index += 1 cur_lr = adjust_learning_rate(cur_lr, optimizer, args.gamma, step_index) if iteration > args.max_iter: break if args.evaluation and iteration % args.eval_step == 0: try: if pEval is None: print("Evaluation started at iteration {} on IC15...". format(iteration)) eval_cmd = "CUDA_VISIBLE_DEVICES=" + str(args.eval_device) + \ " python eval.py" + \ " --tune_from=" + args.save_folder + 'ckpt_' + repr(iteration) + '.pth' + \ " --input_size=1024" + \ " --output_zip=result_temp1" pEval = Popen(eval_cmd, shell=True, stdout=PIPE, stderr=PIPE) elif pEval.poll() is not None: (scorestring, stderrdata) = pEval.communicate() hmean = float( str(scorestring).strip().split(":")[3].split(",") [0].split("}")[0].strip()) writer.add_scalar('test_hmean', hmean, iteration) print("test_hmean for {}-th iter : {:.4f}".format( iteration, hmean)) if pEval is not None: pEval.kill() pEval = None except Exception as e: print("exception happened in evaluation ", e) if pEval is not None: pEval.kill() pEval = None iteration += 1
# Model net = RetinaNet() net.load_state_dict(torch.load('./model/net.pth')) if args.resume: print('==> Resuming from checkpoint..') checkpoint = torch.load('./checkpoint/ckptAdam.pth') net.load_state_dict(checkpoint['net']) best_loss = checkpoint['loss'] start_epoch = checkpoint['epoch'] net = torch.nn.DataParallel(net, device_ids=range(torch.cuda.device_count())) net.cuda() criterion = FocalLoss() # optimizer = optim.SGD(net.parameters(), lr=args.lr, momentum=0.9, weight_decay=1e-4) optimizer = torch.optim.Adam(net.parameters(), lr=args.lr, betas=(0.9, 0.999)) # Training def train(epoch): print('\nEpoch: %d' % epoch) net.train() net.module.freeze_bn() train_loss = 0 for batch_idx, (inputs, loc_targets, cls_targets) in enumerate(trainloader): inputs = Variable(inputs.cuda()) loc_targets = Variable(loc_targets.cuda()) cls_targets = Variable(cls_targets.cuda()) optimizer.zero_grad()
# Model net = RetinaNet() net.load_state_dict(torch.load('./model/net.pth')) if args.resume: print('==> Resuming from checkpoint..') checkpoint = torch.load('./checkpoint/ckpt.pth') net.load_state_dict(checkpoint['net']) best_loss = checkpoint['loss'] start_epoch = checkpoint['epoch'] net = torch.nn.DataParallel(net, device_ids=range(torch.cuda.device_count())) net.cuda() criterion = FocalLoss() optimizer = optim.SGD(net.parameters(), lr=args.lr, momentum=0.9, weight_decay=1e-4) # Training def train(epoch): print('\nEpoch: %d' % epoch) net.train() net.module.freeze_bn() train_loss = 0 for batch_idx, (inputs, loc_targets, cls_targets) in enumerate(trainloader): inputs = Variable(inputs.cuda()) loc_targets = Variable(loc_targets.cuda()) cls_targets = Variable(cls_targets.cuda()) optimizer.zero_grad() loc_preds, cls_preds = net(inputs) loss = criterion(loc_preds, loc_targets, cls_preds, cls_targets)
print("cur_lr : ", cur_lr) print("step_index : ", step_index) print("gpu available : ", torch.cuda.is_available()) print("num_gpus : ", torch.cuda.device_count()) # Set data parallel training net = torch.nn.DataParallel(net, device_ids=[0,1,2,3]) net.cuda() # Training print("==>training start...") net.train() # Freeze BN layer for pre-trained backbone net.module.freeze_bn() # Set optimizer -- SGD or Adam optimizer = optim.SGD(net.parameters(), lr=cur_lr, momentum=0.9, weight_decay=1e-4) #optim.Adam(net.parameters(), lr=cur_lr) # Encode anchor to each feature maps encoder = DataEncoder(cls_thresh=0.5, nms_thresh=0.2) # Tensorboard visualize recorder writer = SummaryWriter(logdir=args.logdir) lossest = 1 save_lossest = False t0 = time.time() for epoch in range(start_epoch, 10000): if iteration > args.max_iter: break for inputs, loc_targets, cls_targets in trainloader: # prepare data and cls & loc label inputs = Variable(inputs.cuda())