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())
