def train():
if not os.path.exists(args.save_folder):
os.mkdir(args.save_folder)
dataset = COCODetection(image_path=cfg.dataset.train_images,
info_file=cfg.dataset.train_info,
transform=SSDAugmentation(MEANS))
if args.validation_epoch > 0:
setup_eval()
val_dataset = COCODetection(image_path=cfg.dataset.valid_images,
info_file=cfg.dataset.valid_info,
transform=BaseTransform(MEANS))
# Parallel wraps the underlying module, but when saving and loading we don't want that
yolact_net = Yolact()
net = yolact_net
net.train()
if args.log:
log = Log(cfg.name,
args.log_folder,
dict(args._get_kwargs()),
overwrite=(args.resume is None),
log_gpu_stats=args.log_gpu)
# I don't use the timer during training (I use a different timing method).
# Apparently there's a race condition with multiple GPUs, so disable it just to be safe.
timer.disable_all()
# Both of these can set args.resume to None, so do them before the check
if args.resume == 'interrupt':
args.resume = SavePath.get_interrupt(args.save_folder)
elif args.resume == 'latest':
args.resume = SavePath.get_latest(args.save_folder, cfg.name)
if args.resume is not None:
print('Resuming training, loading {}...'.format(args.resume))
yolact_net.load_weights(args.resume)
if args.start_iter == -1:
args.start_iter = SavePath.from_str(args.resume).iteration
else:
print('Initializing weights...')
yolact_net.init_weights(backbone_path=args.save_folder +
cfg.backbone.path)
optimizer = optim.SGD(net.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.decay)
criterion = MultiBoxLoss(num_classes=cfg.num_classes,
pos_threshold=cfg.positive_iou_threshold,
neg_threshold=cfg.negative_iou_threshold,
negpos_ratio=cfg.ohem_negpos_ratio)
if args.batch_alloc is not None:
args.batch_alloc = [int(x) for x in args.batch_alloc.split(',')]
if sum(args.batch_alloc) != args.batch_size:
print(
'Error: Batch allocation (%s) does not sum to batch size (%s).'
% (args.batch_alloc, args.batch_size))
exit(-1)
net = CustomDataParallel(NetLoss(net, criterion))
if args.cuda:
net = net.cuda()
# Initialize everything
if not cfg.freeze_bn:
yolact_net.freeze_bn() # Freeze bn so we don't kill our means
yolact_net(torch.zeros(1, 3, cfg.max_size, cfg.max_size).cuda())
if not cfg.freeze_bn: yolact_net.freeze_bn(True)
# loss counters
loc_loss = 0
conf_loss = 0
iteration = max(args.start_iter, 0)
last_time = time.time()
epoch_size = len(dataset) // args.batch_size
num_epochs = math.ceil(cfg.max_iter / epoch_size)
# Which learning rate adjustment step are we on? lr' = lr * gamma ^ step_index
step_index = 0
data_loader = data.DataLoader(dataset,
args.batch_size,
num_workers=args.num_workers,
shuffle=True,
collate_fn=detection_collate,
pin_memory=True)
save_path = lambda epoch, iteration: SavePath(
cfg.name, epoch, iteration).get_path(root=args.save_folder)
time_avg = MovingAverage()
global loss_types # Forms the print order
loss_avgs = {k: MovingAverage(100) for k in loss_types}
print('Begin training!')
print()
# try-except so you can use ctrl+c to save early and stop training
try:
for epoch in range(num_epochs):
# Resume from start_iter
if (epoch + 1) * epoch_size < iteration:
continue
for datum in data_loader:
# Stop if we've reached an epoch if we're resuming from start_iter
if iteration == (epoch + 1) * epoch_size:
break
# Stop at the configured number of iterations even if mid-epoch
if iteration == cfg.max_iter:
break
# Change a config setting if we've reached the specified iteration
changed = False
for change in cfg.delayed_settings:
if iteration >= change[0]:
changed = True
cfg.replace(change[1])
# Reset the loss averages because things might have changed
for avg in loss_avgs:
avg.reset()
# If a config setting was changed, remove it from the list so we don't keep checking
if changed:
cfg.delayed_settings = [
x for x in cfg.delayed_settings if x[0] > iteration
]
# Warm up by linearly interpolating the learning rate from some smaller value
if cfg.lr_warmup_until > 0 and iteration <= cfg.lr_warmup_until:
set_lr(optimizer, (args.lr - cfg.lr_warmup_init) *
(iteration / cfg.lr_warmup_until) +
cfg.lr_warmup_init)
# Adjust the learning rate at the given iterations, but also if we resume from past that iteration
while step_index < len(
cfg.lr_steps
) and iteration >= cfg.lr_steps[step_index]:
step_index += 1
set_lr(optimizer, args.lr * (args.gamma**step_index))
# Zero the grad to get ready to compute gradients
optimizer.zero_grad()
# Forward Pass + Compute loss at the same time (see CustomDataParallel and NetLoss)
losses = net(datum)
losses = {k: (v).mean()
for k, v in losses.items()
} # Mean here because Dataparallel
loss = sum([losses[k] for k in losses])
# no_inf_mean removes some components from the loss, so make sure to backward through all of it
# all_loss = sum([v.mean() for v in losses.values()])
# Backprop
loss.backward(
) # Do this to free up vram even if loss is not finite
if torch.isfinite(loss).item():
optimizer.step()
# Add the loss to the moving average for bookkeeping
for k in losses:
loss_avgs[k].add(losses[k].item())
cur_time = time.time()
elapsed = cur_time - last_time
last_time = cur_time
# Exclude graph setup from the timing information
if iteration != args.start_iter:
time_avg.add(elapsed)
if iteration % 10 == 0:
eta_str = str(
datetime.timedelta(seconds=(cfg.max_iter - iteration) *
time_avg.get_avg())).split('.')[0]
total = sum([loss_avgs[k].get_avg() for k in losses])
loss_labels = sum([[k, loss_avgs[k].get_avg()]
for k in loss_types if k in losses], [])
print(('[%3d] %7d ||' + (' %s: %.3f |' * len(losses)) +
' T: %.3f || ETA: %s || timer: %.3f') %
tuple([epoch, iteration] + loss_labels +
[total, eta_str, elapsed]),
flush=True)
if args.log:
precision = 5
loss_info = {
k: round(losses[k].item(), precision)
for k in losses
}
loss_info['T'] = round(loss.item(), precision)
if args.log_gpu:
log.log_gpu_stats = (iteration % 10 == 0
) # nvidia-smi is sloooow
log.log('train',
loss=loss_info,
epoch=epoch,
iter=iteration,
lr=round(cur_lr, 10),
elapsed=elapsed)
log.log_gpu_stats = args.log_gpu
iteration += 1
if iteration % args.save_interval == 0 and iteration != args.start_iter:
if args.keep_latest:
latest = SavePath.get_latest(args.save_folder,
cfg.name)
print('Saving state, iter:', iteration)
yolact_net.save_weights(save_path(epoch, iteration))
if args.keep_latest and latest is not None:
if args.keep_latest_interval <= 0 or iteration % args.keep_latest_interval != args.save_interval:
print('Deleting old save...')
os.remove(latest)
# This is done per epoch
if args.validation_epoch > 0:
if epoch % args.validation_epoch == 0 and epoch > 0:
compute_validation_map(epoch, iteration, yolact_net,
val_dataset,
log if args.log else None)
# Compute validation mAP after training is finished
compute_validation_map(epoch, iteration, yolact_net, val_dataset,
log if args.log else None)
except KeyboardInterrupt:
if args.interrupt:
print('Stopping early. Saving network...')
# Delete previous copy of the interrupted network so we don't spam the weights folder
SavePath.remove_interrupt(args.save_folder)
yolact_net.save_weights(
save_path(epoch,
repr(iteration) + '_interrupt'))
exit()
yolact_net.save_weights(save_path(epoch, iteration))
def train():
if not os.path.exists(args.save_folder):
os.mkdir(args.save_folder)
dataset = COCODetection(image_path=cfg.dataset.train_images,
info_file=cfg.dataset.train_info,
transform=SSDAugmentation(MEANS))
if args.validation_epoch > 0:
setup_eval()
val_dataset = COCODetection(image_path=cfg.dataset.valid_images,
info_file=cfg.dataset.valid_info,
transform=BaseTransform(MEANS))
# Parallel wraps the underlying module, but when saving and loading we don't want that
yolact_net = Yolact()
net = yolact_net
net.train()
# I don't use the timer during training (I use a different timing method).
# Apparently there's a race condition with multiple GPUs.
timer.disable_all()
# Both of these can set args.resume to None, so do them before the check
if args.resume == 'interrupt':
args.resume = SavePath.get_interrupt(args.save_folder)
elif args.resume == 'latest':
args.resume = SavePath.get_latest(args.save_folder, cfg.name)
if args.resume is not None:
print('Resuming training, loading {}...'.format(args.resume))
yolact_net.load_weights(args.resume)
if args.start_iter == -1:
args.start_iter = SavePath.from_str(args.resume).iteration
else:
print('Initializing weights...')
yolact_net.init_weights(backbone_path=args.save_folder + cfg.backbone.path)
optimizer = optim.SGD(net.parameters(), lr=args.lr, momentum=args.momentum,
weight_decay=args.decay)
criterion = MultiBoxLoss(num_classes=cfg.num_classes,
pos_threshold=cfg.positive_iou_threshold,
neg_threshold=cfg.negative_iou_threshold,
negpos_ratio=3)
if args.cuda:
cudnn.benchmark = True
net = nn.DataParallel(net).cuda()
criterion = nn.DataParallel(criterion).cuda()
# net = net.cuda()
# criterion = criterion.cuda()
# criterion = criterion.cuda()
# loss counters
loc_loss = 0
conf_loss = 0
iteration = max(args.start_iter, 0)
last_time = time.time()
epoch_size = len(dataset) // args.batch_size
num_epochs = math.ceil(cfg.max_iter / epoch_size)
# Which learning rate adjustment step are we on? lr' = lr * gamma ^ step_index
step_index = 0
data_loader = data.DataLoader(dataset, args.batch_size,
num_workers=args.num_workers,
shuffle=True, collate_fn=detection_collate,
pin_memory=True)
save_path = lambda epoch, iteration: SavePath(cfg.name, epoch, iteration).get_path(root=args.save_folder)
time_avg = MovingAverage()
global loss_types # Forms the print order
loss_avgs = {k: MovingAverage(100) for k in loss_types}
print('Begin training!')
print()
# try-except so you can use ctrl+c to save early and stop training
try:
for epoch in range(num_epochs):
# Resume from start_iter
if (epoch + 1) * epoch_size < iteration:
continue
for datum in data_loader:
# Stop if we've reached an epoch if we're resuming from start_iter
if iteration == (epoch + 1) * epoch_size:
break
# Stop at the configured number of iterations even if mid-epoch
if iteration == cfg.max_iter:
break
# Change a config setting if we've reached the specified iteration
changed = False
for change in cfg.delayed_settings:
if iteration >= change[0]:
changed = True
cfg.replace(change[1])
# Reset the loss averages because things might have changed
for avg in loss_avgs:
avg.reset()
# If a config setting was changed, remove it from the list so we don't keep checking
if changed:
cfg.delayed_settings = [x for x in cfg.delayed_settings if x[0] > iteration]
# Warm up by linearly interpolating the learning rate from some smaller value
if cfg.lr_warmup_until > 0 and iteration <= cfg.lr_warmup_until:
set_lr(optimizer,
(args.lr - cfg.lr_warmup_init) * (iteration / cfg.lr_warmup_until) + cfg.lr_warmup_init)
# Adjust the learning rate at the given iterations, but also if we resume from past that iteration
while step_index < len(cfg.lr_steps) and iteration >= cfg.lr_steps[step_index]:
step_index += 1
set_lr(optimizer, args.lr * (args.gamma ** step_index))
# Load training data
# Note, for training on multiple gpus this will use the custom replicate and gather I wrote up there
images, targets, masks, num_crowds = prepare_data(datum)
# Forward Pass
out = net(images)
# Compute Loss
optimizer.zero_grad()
wrapper = ScatterWrapper(targets, masks, num_crowds)
losses = criterion(out, wrapper, wrapper.make_mask())
losses = {k: v.mean() for k, v in losses.items()} # Mean here because Dataparallel
# original
loss = sum([losses[k] for k in losses])
# loss = sum([losses[k] for k in losses]) + losses['S'] * 10 # Huan
# Backprop
loss.backward() # Do this to free up vram even if loss is not finite
if torch.isfinite(loss).item():
optimizer.step()
# Add the loss to the moving average for bookkeeping
for k in losses:
loss_avgs[k].add(losses[k].item())
cur_time = time.time()
elapsed = cur_time - last_time
last_time = cur_time
# Exclude graph setup from the timing information
if iteration != args.start_iter:
time_avg.add(elapsed)
if iteration % 100 == 0:
eta_str = \
str(datetime.timedelta(seconds=(cfg.max_iter - iteration) * time_avg.get_avg())).split('.')[0]
total = sum([loss_avgs[k].get_avg() for k in losses])
loss_labels = sum([[k, loss_avgs[k].get_avg()] for k in loss_types if k in losses], [])
print(('[%3d] %7d ||' + (' %s: %.3f |' * len(losses)) + ' T: %.3f || ETA: %s || timer: %.3f')
% tuple([epoch, iteration] + loss_labels + [total, eta_str, elapsed]), flush=True)
iteration += 1
if iteration % args.save_interval == 0 and iteration != args.start_iter:
if args.keep_latest:
latest = SavePath.get_latest(args.save_folder, cfg.name)
print('Saving state, iter:', iteration)
yolact_net.save_weights(save_path(epoch, iteration))
if args.keep_latest and latest is not None:
if args.keep_latest_interval <= 0 or iteration % args.keep_latest_interval != args.save_interval:
print('Deleting old save...')
os.remove(latest)
# This is done per epoch
if args.validation_epoch > 0:
if epoch % args.validation_epoch == 0 and epoch > 0:
compute_validation_map(yolact_net, val_dataset)
except KeyboardInterrupt:
print('Stopping early. Saving network...')
# Delete previous copy of the interrupted network so we don't spam the weights folder
SavePath.remove_interrupt(args.save_folder)
yolact_net.save_weights(save_path(epoch, repr(iteration) + '_interrupt'))
exit()
yolact_net.save_weights(save_path(epoch, iteration))
def train():
if not os.path.exists(args.save_folder):
os.mkdir(args.save_folder)
dataset = COCODetection(image_path=cfg.dataset.train_images,
info_file=cfg.dataset.train_info,
transform=SSDAugmentation(MEANS))
if args.validation_epoch > 0:
setup_eval()
val_dataset = COCODetection(image_path=cfg.dataset.valid_images,
info_file=cfg.dataset.valid_info,
transform=BaseTransform(MEANS))
# Parallel wraps the underlying module, but when saving and loading we don't want that
yolact_net = Yolact()
net = yolact_net
net.train()
print('\n--- Generator created! ---')
# NOTE
# I maunally set the original image size and seg size as 138
# might change in the future, for example 550
if cfg.pred_seg:
dis_size = 138
dis_net = Discriminator_Dcgan(i_size=dis_size, s_size=dis_size)
# set the dis net's initial parameter values
dis_net.apply(gan_init)
dis_net.train()
print('--- Discriminator created! ---\n')
if args.log:
log = Log(cfg.name,
args.log_folder,
dict(args._get_kwargs()),
overwrite=(args.resume is None),
log_gpu_stats=args.log_gpu)
# I don't use the timer during training (I use a different timing method).
# Apparently there's a race condition with multiple GPUs, so disable it just to be safe.
timer.disable_all()
# Both of these can set args.resume to None, so do them before the check
if args.resume == 'interrupt':
args.resume = SavePath.get_interrupt(args.save_folder)
elif args.resume == 'latest':
args.resume = SavePath.get_latest(args.save_folder, cfg.name)
if args.resume is not None:
print('Resuming training, loading {}...'.format(args.resume))
yolact_net.load_weights(args.resume)
if args.start_iter == -1:
args.start_iter = SavePath.from_str(args.resume).iteration
else:
print('Initializing weights...')
yolact_net.init_weights(backbone_path=args.save_folder +
cfg.backbone.path)
# optimizer_gen = optim.SGD(net.parameters(), lr=args.lr, momentum=args.momentum,
# weight_decay=args.decay)
# if cfg.pred_seg:
# optimizer_dis = optim.SGD(dis_net.parameters(), lr=cfg.dis_lr, momentum=args.momentum,
# weight_decay=args.decay)
# schedule_dis = ReduceLROnPlateau(optimizer_dis, mode = 'min', patience=6, min_lr=1E-6)
# NOTE: Using the Ranger Optimizer for the generator
optimizer_gen = Ranger(net.parameters(),
lr=args.lr,
weight_decay=args.decay)
if cfg.pred_seg:
optimizer_dis = optim.SGD(dis_net.parameters(), lr=cfg.dis_lr)
schedule_dis = ReduceLROnPlateau(optimizer_dis,
mode='min',
patience=6,
min_lr=1E-6)
criterion = MultiBoxLoss(num_classes=cfg.num_classes,
pos_threshold=cfg.positive_iou_threshold,
neg_threshold=cfg.negative_iou_threshold,
negpos_ratio=cfg.ohem_negpos_ratio,
pred_seg=cfg.pred_seg)
criterion_dis = nn.BCELoss()
if args.batch_alloc is not None:
args.batch_alloc = [int(x) for x in args.batch_alloc.split(',')]
if sum(args.batch_alloc) != args.batch_size:
print(
'Error: Batch allocation (%s) does not sum to batch size (%s).'
% (args.batch_alloc, args.batch_size))
exit(-1)
net = CustomDataParallel(NetLoss(net, criterion, pred_seg=cfg.pred_seg))
if args.cuda:
net = net.cuda()
if cfg.pred_seg:
dis_net = dis_net.cuda()
# Initialize everything
if not cfg.freeze_bn:
yolact_net.freeze_bn() # Freeze bn so we don't kill our means
yolact_net(torch.zeros(1, 3, cfg.max_size, cfg.max_size).cuda())
if not cfg.freeze_bn: yolact_net.freeze_bn(True)
# loss counters
loc_loss = 0
conf_loss = 0
iteration = max(args.start_iter, 0)
last_time = time.time()
epoch_size = len(dataset) // args.batch_size
num_epochs = math.ceil(cfg.max_iter / epoch_size)
# Which learning rate adjustment step are we on? lr' = lr * gamma ^ step_index
step_index = 0
data_loader = data.DataLoader(dataset,
args.batch_size,
num_workers=args.num_workers,
shuffle=True,
collate_fn=detection_collate,
pin_memory=True)
# NOTE
val_loader = data.DataLoader(val_dataset,
args.batch_size,
num_workers=args.num_workers * 2,
shuffle=True,
collate_fn=detection_collate,
pin_memory=True)
save_path = lambda epoch, iteration: SavePath(
cfg.name, epoch, iteration).get_path(root=args.save_folder)
time_avg = MovingAverage()
global loss_types # Forms the print order
# TODO: global command can modify global variable inside of the function.
loss_avgs = {k: MovingAverage(100) for k in loss_types}
print('Begin training!')
print()
# try-except so you can use ctrl+c to save early and stop training
try:
for epoch in range(num_epochs):
# Resume from start_iter
if (epoch + 1) * epoch_size < iteration:
continue
for datum in data_loader:
# Stop if we've reached an epoch if we're resuming from start_iter
if iteration == (epoch + 1) * epoch_size:
break
# Stop at the configured number of iterations even if mid-epoch
if iteration == cfg.max_iter:
break
# Change a config setting if we've reached the specified iteration
changed = False
for change in cfg.delayed_settings:
if iteration >= change[0]:
changed = True
cfg.replace(change[1])
# Reset the loss averages because things might have changed
for avg in loss_avgs:
avg.reset()
# If a config setting was changed, remove it from the list so we don't keep checking
if changed:
cfg.delayed_settings = [
x for x in cfg.delayed_settings if x[0] > iteration
]
# Warm up by linearly interpolating the learning rate from some smaller value
if cfg.lr_warmup_until > 0 and iteration <= cfg.lr_warmup_until:
set_lr(optimizer_gen, (args.lr - cfg.lr_warmup_init) *
(iteration / cfg.lr_warmup_until) +
cfg.lr_warmup_init)
# Adjust the learning rate at the given iterations, but also if we resume from past that iteration
while step_index < len(
cfg.lr_steps
) and iteration >= cfg.lr_steps[step_index]:
step_index += 1
set_lr(optimizer_gen, args.lr * (args.gamma**step_index))
# Zero the grad to get ready to compute gradients
optimizer_gen.zero_grad()
# NOTE
if cfg.pred_seg:
# ====== GAN Train ======
optimizer_dis.zero_grad()
dis_net.zero_grad()
# train the gen and dis in different iteration count
it_alter_period = iteration % (cfg.gen_iter + cfg.dis_iter)
if it_alter_period >= cfg.gen_iter:
freeze_pretrain(yolact_net, freeze=True)
freeze_pretrain(net, freeze=True)
freeze_pretrain(dis_net, freeze=False)
if it_alter_period == (cfg.gen_iter + 1):
print('--- Generator freeze ---')
print('--- Discriminator training ---')
# ----- Discriminator part -----
# seg_list is the prediction mask
# can be regarded as generated images from YOLACT
# pred_list is the prediction label
# seg_list dim: list of (138,138,instances)
# pred_list dim: list of (instances)
losses, seg_list, pred_list = net(datum)
seg_clas, mask_clas, b, seg_size = seg_mask_clas(
seg_list, pred_list, datum)
# input image size is [b, 3, 550, 550]
# downsample to [b, 3, seg_h, seg_w]
image = interpolate(torch.stack(datum[0]),
size=seg_size,
mode='bilinear',
align_corners=False)
# Because in the discriminator training, we do not
# want the gradient flow back to the generator part
# we detach seg_clas (mask_clas come the data, does not have grad)
seg_input = seg_clas.clone().detach()
output_pred = dis_net(img=image, seg=seg_input)
output_grou = dis_net(img=image, seg=mask_clas)
if iteration % (cfg.gen_iter + cfg.dis_iter) == 0:
print(
f'Probability of fake is fake: {output_pred.mean().item():.2f}'
)
print(
f'Probability of real is real: {output_grou.mean().item():.2f}'
)
# 0 for Fake/Generated
# 1 for True/Ground Truth
fake_label = torch.zeros(b)
real_label = torch.ones(b)
# Advice of practical implementation
# from https://arxiv.org/abs/1611.08408
# loss_pred = -criterion_dis(output_pred,target=real_label)
loss_pred = criterion_dis(output_pred,
target=fake_label)
loss_grou = criterion_dis(output_grou,
target=real_label)
loss_dis = loss_pred + loss_grou
# TODO: Grid Search this one
lambda_dis = cfg.lambda_dis
loss_dis = lambda_dis * loss_dis
# Backprop of the discriminator
loss_dis.backward()
optimizer_dis.step()
else:
freeze_pretrain(yolact_net, freeze=False)
freeze_pretrain(net, freeze=False)
freeze_pretrain(dis_net, freeze=True)
if it_alter_period == 0:
print('--- Generator training ---')
print('--- Discriminator freeze ---')
# ----- Generator part -----
net.zero_grad()
losses, seg_list, pred_list = net(datum)
seg_clas, mask_clas, b, seg_size = seg_mask_clas(
seg_list, pred_list, datum)
image = interpolate(torch.stack(datum[0]),
size=seg_size,
mode='bilinear',
align_corners=False)
# Perform forward pass of all-fake batch through D
# NOTE that seg_clas CANNOT detach, in order to flow the
# gradient back to the generator
output = dis_net(img=image, seg=seg_clas)
# Since the log(1-D(G(x))) not provide sufficient gradients
# We want log(D(G(x)) instead, this can be achieve by
# use the real_label as target.
# This step is crucial for the information of discriminator
# to go into the generator.
# Calculate G's loss based on this output
real_label = torch.ones(b)
loss_gen = criterion_dis(output, target=real_label)
if not it_alter_period >= cfg.gen_iter:
# since the dis is already freeze, the gradients will only
# record the YOLACT
loss_gen.backward()
# Do this to free up vram even if loss is not finite
losses = {
k: (v).mean()
for k, v in losses.items()
} # Mean here because Dataparallel
loss = sum([losses[k] for k in losses])
if torch.isfinite(loss).item():
# since the optimizer_gen is for YOLACT only
# only the gen will be updated
optimizer_gen.step()
else:
# ====== Normal YOLACT Train ======
# Forward Pass + Compute loss at the same time (see CustomDataParallel and NetLoss)
losses = net(datum)
losses = {k: (v).mean()
for k, v in losses.items()
} # Mean here because Dataparallel
loss = sum([losses[k] for k in losses])
# no_inf_mean removes some components from the loss, so make sure to backward through all of it
# all_loss = sum([v.mean() for v in losses.values()])
# Backprop
loss.backward(
) # Do this to free up vram even if loss is not finite
if torch.isfinite(loss).item():
optimizer_gen.step()
# Add the loss to the moving average for bookkeeping
for k in losses:
loss_avgs[k].add(losses[k].item())
cur_time = time.time()
elapsed = cur_time - last_time
last_time = cur_time
# Exclude graph setup from the timing information
if iteration != args.start_iter:
time_avg.add(elapsed)
if iteration % 10 == 0:
eta_str = str(
datetime.timedelta(seconds=(cfg.max_iter - iteration) *
time_avg.get_avg())).split('.')[0]
total = sum([loss_avgs[k].get_avg() for k in losses])
loss_labels = sum([[k, loss_avgs[k].get_avg()]
for k in loss_types if k in losses], [])
if cfg.pred_seg and (it_alter_period >= cfg.gen_iter):
print(
('[%3d] %7d ||' + (' %s: %.3f |' * len(losses)) +
' T: %.3f || Dis: %.2f || ETA: %s || timer: %.3f')
% tuple([epoch, iteration] + loss_labels +
[total, loss_dis, eta_str, elapsed]),
flush=True)
else:
print(('[%3d] %7d ||' + (' %s: %.3f |' * len(losses)) +
' T: %.3f || ETA: %s || timer: %.3f') %
tuple([epoch, iteration] + loss_labels +
[total, eta_str, elapsed]),
flush=True)
# Loss Key:
# - B: Box Localization Loss
# - C: Class Confidence Loss
# - M: Mask Loss
# - P: Prototype Loss
# - D: Coefficient Diversity Loss
# - E: Class Existence Loss
# - S: Semantic Segmentation Loss
# - T: Total loss
# -Dis:Discriminator Loss
if args.log:
precision = 5
loss_info = {
k: round(losses[k].item(), precision)
for k in losses
}
loss_info['T'] = round(loss.item(), precision)
if args.log_gpu:
log.log_gpu_stats = (iteration % 10 == 0
) # nvidia-smi is sloooow
log.log('train',
loss=loss_info,
epoch=epoch,
iter=iteration,
lr=round(cur_lr, 10),
elapsed=elapsed)
log.log_gpu_stats = args.log_gpu
iteration += 1
if iteration % args.save_interval == 0 and iteration != args.start_iter:
if args.keep_latest:
latest = SavePath.get_latest(args.save_folder,
cfg.name)
print('Saving state, iter:', iteration)
yolact_net.save_weights(save_path(epoch, iteration))
if args.keep_latest and latest is not None:
if args.keep_latest_interval <= 0 or iteration % args.keep_latest_interval != args.save_interval:
print('Deleting old save...')
os.remove(latest)
# This is done per epoch
if args.validation_epoch > 0:
# NOTE: Validation loss
# if cfg.pred_seg:
# net.eval()
# dis_net.eval()
# cfg.gan_eval = True
# with torch.no_grad():
# for datum in tqdm(val_loader, desc='GAN Validation'):
# losses, seg_list, pred_list = net(datum)
# losses, seg_list, pred_list = net(datum)
# # TODO: warp below as a function
# seg_list = [v.permute(2,1,0).contiguous() for v in seg_list]
# b = len(seg_list) # batch size
# _, seg_h, seg_w = seg_list[0].size()
# seg_clas = torch.zeros(b, cfg.num_classes-1, seg_h, seg_w)
# mask_clas = torch.zeros(b, cfg.num_classes-1, seg_h, seg_w)
# target_list = [target for target in datum[1][0]]
# mask_list = [interpolate(mask.unsqueeze(0), size = (seg_h,seg_w),mode='bilinear', \
# align_corners=False).squeeze() for mask in datum[1][1]]
# for idx in range(b):
# for i, (pred, i_target) in enumerate(zip(pred_list[idx], target_list[idx])):
# seg_clas[idx, pred, ...] += seg_list[idx][i,...]
# mask_clas[idx, i_target[-1].long(), ...] += mask_list[idx][i,...]
# seg_clas = torch.clamp(seg_clas, 0, 1)
# image = interpolate(torch.stack(datum[0]), size = (seg_h,seg_w),
# mode='bilinear',align_corners=False)
# real_label = torch.ones(b)
# output_pred = dis_net(img = image, seg = seg_clas)
# output_grou = dis_net(img = image, seg = mask_clas)
# loss_pred = -criterion_dis(output_pred,target=real_label)
# loss_grou = criterion_dis(output_grou,target=real_label)
# loss_dis = loss_pred + loss_grou
# losses = { k: (v).mean() for k,v in losses.items() }
# loss = sum([losses[k] for k in losses])
# val_loss = loss - cfg.lambda_dis*loss_dis
# schedule_dis.step(loss_dis)
# lr = [group['lr'] for group in optimizer_dis.param_groups]
# print(f'Discriminator lr: {lr[0]}')
# net.train()
if epoch % args.validation_epoch == 0 and epoch > 0:
cfg.gan_eval = False
dis_net.eval()
compute_validation_map(epoch, iteration, yolact_net,
val_dataset,
log if args.log else None)
# Compute validation mAP after training is finished
compute_validation_map(epoch, iteration, yolact_net, val_dataset,
log if args.log else None)
except KeyboardInterrupt:
if args.interrupt:
print('Stopping early. Saving network...')
# Delete previous copy of the interrupted network so we don't spam the weights folder
SavePath.remove_interrupt(args.save_folder)
yolact_net.save_weights(
save_path(epoch,
repr(iteration) + '_interrupt'))
exit()
yolact_net.save_weights(save_path(epoch, iteration))
def train(rank, args):
if args.num_gpus > 1:
multi_gpu_rescale(args)
if rank == 0:
if not os.path.exists(args.save_folder):
os.mkdir(args.save_folder)
# fix the seed for reproducibility
seed = args.random_seed + rank
torch.manual_seed(seed)
np.random.seed(seed)
random.seed(seed)
# set up logger
setup_logger(output=os.path.join(args.log_folder, cfg.name),
distributed_rank=rank)
logger = logging.getLogger("yolact.train")
w = SummaryHelper(distributed_rank=rank,
log_dir=os.path.join(args.log_folder, cfg.name))
w.add_text("argv", " ".join(sys.argv))
logger.info("Args: {}".format(" ".join(sys.argv)))
import git
with git.Repo(search_parent_directories=True) as repo:
w.add_text("git_hash", repo.head.object.hexsha)
logger.info("git hash: {}".format(repo.head.object.hexsha))
if args.num_gpus > 1:
try:
logger.info("Initializing torch.distributed backend...")
dist.init_process_group(backend='nccl',
init_method=args.dist_url,
world_size=args.num_gpus,
rank=rank)
except Exception as e:
logger.error("Process group URL: {}".format(args.dist_url))
raise e
misc.barrier()
if torch.cuda.device_count() > 1:
logger.info('Multiple GPUs detected! Turning off JIT.')
collate_fn = detection_collate
if cfg.dataset.name == 'YouTube VIS':
dataset = YoutubeVIS(image_path=cfg.dataset.train_images,
info_file=cfg.dataset.train_info,
configs=cfg.dataset,
transform=SSDAugmentationVideo(MEANS))
if cfg.dataset.joint == 'coco':
joint_dataset = COCODetection(
image_path=cfg.joint_dataset.train_images,
info_file=cfg.joint_dataset.train_info,
transform=SSDAugmentation(MEANS))
joint_collate_fn = detection_collate
if args.validation_epoch > 0:
setup_eval()
val_dataset = YoutubeVIS(image_path=cfg.dataset.valid_images,
info_file=cfg.dataset.valid_info,
configs=cfg.dataset,
transform=BaseTransformVideo(MEANS))
collate_fn = collate_fn_youtube_vis
elif cfg.dataset.name == 'FlyingChairs':
dataset = FlyingChairs(image_path=cfg.dataset.trainval_images,
info_file=cfg.dataset.trainval_info)
collate_fn = collate_fn_flying_chairs
else:
dataset = COCODetection(image_path=cfg.dataset.train_images,
info_file=cfg.dataset.train_info,
transform=SSDAugmentation(MEANS))
if args.validation_epoch > 0:
setup_eval()
val_dataset = COCODetection(image_path=cfg.dataset.valid_images,
info_file=cfg.dataset.valid_info,
transform=BaseTransform(MEANS))
# Set cuda device early to avoid duplicate model in master GPU
if args.cuda:
torch.cuda.set_device(rank)
# Parallel wraps the underlying module, but when saving and loading we don't want that
yolact_net = Yolact()
net = yolact_net
net.train()
# I don't use the timer during training (I use a different timing method).
# Apparently there's a race condition with multiple GPUs.
# use timer for experiments
timer.disable_all()
# Both of these can set args.resume to None, so do them before the check
if args.resume == 'interrupt':
args.resume = SavePath.get_interrupt(args.save_folder)
elif args.resume == 'latest':
args.resume = SavePath.get_latest(args.save_folder, cfg.name)
if args.resume is not None:
logger.info('Resuming training, loading {}...'.format(args.resume))
yolact_net.load_weights(args.resume, args=args)
if args.start_iter == -1:
args.start_iter = SavePath.from_str(args.resume).iteration
else:
logger.info('Initializing weights...')
yolact_net.init_weights(backbone_path=args.save_folder +
cfg.backbone.path)
if cfg.flow.train_flow:
criterion = OpticalFlowLoss()
else:
criterion = MultiBoxLoss(num_classes=cfg.num_classes,
pos_threshold=cfg.positive_iou_threshold,
neg_threshold=cfg.negative_iou_threshold,
negpos_ratio=3)
if args.cuda:
net.cuda(rank)
if misc.is_distributed_initialized():
net = nn.parallel.DistributedDataParallel(
net,
device_ids=[rank],
output_device=rank,
broadcast_buffers=False,
find_unused_parameters=True)
optimizer = optim.SGD(filter(lambda x: x.requires_grad, net.parameters()),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.decay)
# loss counters
iteration = max(args.start_iter, 0)
w.set_step(iteration)
last_time = time.time()
epoch_size = len(dataset) // args.batch_size // args.num_gpus
num_epochs = math.ceil(cfg.max_iter / epoch_size)
# Which learning rate adjustment step are we on? lr' = lr * gamma ^ step_index
step_index = 0
from data.sampler_utils import InfiniteSampler, build_batch_data_sampler
infinite_sampler = InfiniteSampler(dataset,
seed=args.random_seed,
num_replicas=args.num_gpus,
rank=rank,
shuffle=True)
train_sampler = build_batch_data_sampler(infinite_sampler,
images_per_batch=args.batch_size)
data_loader = data.DataLoader(
dataset,
num_workers=args.num_workers,
collate_fn=collate_fn,
multiprocessing_context="fork" if args.num_workers > 1 else None,
batch_sampler=train_sampler)
data_loader_iter = iter(data_loader)
if cfg.dataset.joint:
joint_infinite_sampler = InfiniteSampler(joint_dataset,
seed=args.random_seed,
num_replicas=args.num_gpus,
rank=rank,
shuffle=True)
joint_train_sampler = build_batch_data_sampler(
joint_infinite_sampler, images_per_batch=args.batch_size)
joint_data_loader = data.DataLoader(
joint_dataset,
num_workers=args.num_workers,
collate_fn=joint_collate_fn,
multiprocessing_context="fork" if args.num_workers > 1 else None,
batch_sampler=joint_train_sampler)
joint_data_loader_iter = iter(joint_data_loader)
save_path = lambda epoch, iteration: SavePath(
cfg.name, epoch, iteration).get_path(root=args.save_folder)
time_avg = MovingAverage()
data_time_avg = MovingAverage(10)
global loss_types # Forms the print order
loss_avgs = {k: MovingAverage(100) for k in loss_types}
def backward_and_log(prefix,
net_outs,
targets,
masks,
num_crowds,
extra_loss=None):
optimizer.zero_grad()
out = net_outs["pred_outs"]
losses = criterion(out, targets, masks, num_crowds)
losses = {k: v.mean()
for k, v in losses.items()} # Mean here because Dataparallel
if extra_loss is not None:
assert type(extra_loss) == dict
losses.update(extra_loss)
loss = sum([losses[k] for k in losses])
# Backprop
loss.backward() # Do this to free up vram even if loss is not finite
if torch.isfinite(loss).item():
optimizer.step()
# Add the loss to the moving average for bookkeeping
for k in losses:
loss_avgs[k].add(losses[k].item())
w.add_scalar('{prefix}/{key}'.format(prefix=prefix, key=k),
losses[k].item())
return losses
logger.info('Begin training!')
# try-except so you can use ctrl+c to save early and stop training
try:
for epoch in range(num_epochs):
# Resume from start_iter
if (epoch + 1) * epoch_size < iteration:
continue
while True:
data_start_time = time.perf_counter()
datum = next(data_loader_iter)
data_end_time = time.perf_counter()
data_time = data_end_time - data_start_time
if iteration != args.start_iter:
data_time_avg.add(data_time)
# Stop if we've reached an epoch if we're resuming from start_iter
if iteration == (epoch + 1) * epoch_size:
break
# Stop at the configured number of iterations even if mid-epoch
if iteration == cfg.max_iter:
break
# Change a config setting if we've reached the specified iteration
changed = False
for change in cfg.delayed_settings:
if iteration >= change[0]:
changed = True
cfg.replace(change[1])
# Reset the loss averages because things might have changed
for avg in loss_avgs:
avg.reset()
# If a config setting was changed, remove it from the list so we don't keep checking
if changed:
cfg.delayed_settings = [
x for x in cfg.delayed_settings if x[0] > iteration
]
# Warm up by linearly interpolating the learning rate from some smaller value
if cfg.lr_warmup_until > 0 and iteration <= cfg.lr_warmup_until and cfg.lr_warmup_init < args.lr:
set_lr(optimizer, (args.lr - cfg.lr_warmup_init) *
(iteration / cfg.lr_warmup_until) +
cfg.lr_warmup_init)
elif cfg.lr_schedule == 'cosine':
set_lr(
optimizer,
args.lr *
((math.cos(math.pi * iteration / cfg.max_iter) + 1.) *
.5))
# Adjust the learning rate at the given iterations, but also if we resume from past that iteration
while cfg.lr_schedule == 'step' and step_index < len(
cfg.lr_steps
) and iteration >= cfg.lr_steps[step_index]:
step_index += 1
set_lr(optimizer, args.lr * (args.gamma**step_index))
global lr
w.add_scalar('meta/lr', lr)
if cfg.dataset.name == "FlyingChairs":
imgs_1, imgs_2, flows = prepare_flow_data(datum)
net_outs = net(None, extras=(imgs_1, imgs_2))
# Compute Loss
optimizer.zero_grad()
losses = criterion(net_outs, flows)
losses = {k: v.mean()
for k, v in losses.items()
} # Mean here because Dataparallel
loss = sum([losses[k] for k in losses])
# Backprop
loss.backward(
) # Do this to free up vram even if loss is not finite
if torch.isfinite(loss).item():
optimizer.step()
# Add the loss to the moving average for bookkeeping
for k in losses:
loss_avgs[k].add(losses[k].item())
w.add_scalar('loss/%s' % k, losses[k].item())
elif cfg.dataset.joint or not cfg.dataset.is_video:
if cfg.dataset.joint:
joint_datum = next(joint_data_loader_iter)
# Load training data
# Note, for training on multiple gpus this will use the custom replicate and gather I wrote up there
images, targets, masks, num_crowds = prepare_data(
joint_datum)
else:
images, targets, masks, num_crowds = prepare_data(
datum)
extras = {
"backbone": "full",
"interrupt": False,
"moving_statistics": {
"aligned_feats": []
}
}
net_outs = net(images, extras=extras)
run_name = "joint" if cfg.dataset.joint else "compute"
losses = backward_and_log(run_name, net_outs, targets,
masks, num_crowds)
# Forward Pass
if cfg.dataset.is_video:
# reference frames
references = []
moving_statistics = {"aligned_feats": [], "conf_hist": []}
for idx, frame in enumerate(datum[:0:-1]):
images, annots = frame
extras = {
"backbone": "full",
"interrupt": True,
"keep_statistics": True,
"moving_statistics": moving_statistics
}
with torch.no_grad():
net_outs = net(images, extras=extras)
moving_statistics["feats"] = net_outs["feats"]
moving_statistics["lateral"] = net_outs["lateral"]
keys_to_save = ("outs_phase_1", "outs_phase_2")
for key in set(net_outs.keys()) - set(keys_to_save):
del net_outs[key]
references.append(net_outs)
# key frame with annotation, but not compute full backbone
frame = datum[0]
images, annots = frame
frame = (
images,
annots,
)
images, targets, masks, num_crowds = prepare_data(frame)
extras = {
"backbone": "full",
"interrupt": not cfg.flow.base_backward,
"moving_statistics": moving_statistics
}
gt_net_outs = net(images, extras=extras)
if cfg.flow.base_backward:
losses = backward_and_log("compute", gt_net_outs,
targets, masks, num_crowds)
keys_to_save = ("outs_phase_1", "outs_phase_2")
for key in set(gt_net_outs.keys()) - set(keys_to_save):
del gt_net_outs[key]
# now do the warp
if len(references) > 0:
reference_frame = references[0]
extras = {
"backbone": "partial",
"moving_statistics": moving_statistics
}
net_outs = net(images, extras=extras)
extra_loss = yolact_net.extra_loss(
net_outs, gt_net_outs)
losses = backward_and_log("warp",
net_outs,
targets,
masks,
num_crowds,
extra_loss=extra_loss)
cur_time = time.time()
elapsed = cur_time - last_time
last_time = cur_time
w.add_scalar('meta/data_time', data_time)
w.add_scalar('meta/iter_time', elapsed)
# Exclude graph setup from the timing information
if iteration != args.start_iter:
time_avg.add(elapsed)
if iteration % 10 == 0:
eta_str = str(
datetime.timedelta(seconds=(cfg.max_iter - iteration) *
time_avg.get_avg())).split('.')[0]
if torch.cuda.is_available():
max_mem_mb = torch.cuda.max_memory_allocated(
) / 1024.0 / 1024.0
# torch.cuda.reset_max_memory_allocated()
else:
max_mem_mb = None
logger.info("""\
eta: {eta} epoch: {epoch} iter: {iter} \
{losses} {loss_total} \
time: {time} data_time: {data_time} lr: {lr} {memory}\
""".format(eta=eta_str,
epoch=epoch,
iter=iteration,
losses=" ".join([
"{}: {:.3f}".format(k, loss_avgs[k].get_avg()) for k in losses
]),
loss_total="T: {:.3f}".format(
sum([loss_avgs[k].get_avg() for k in losses])),
data_time="{:.3f}".format(data_time_avg.get_avg()),
time="{:.3f}".format(elapsed),
lr="{:.6f}".format(lr),
memory="max_mem: {:.0f}M".format(max_mem_mb)))
if rank == 0 and iteration % 100 == 0:
if cfg.flow.train_flow:
import flowiz as fz
from layers.warp_utils import deform_op
tgt_size = (64, 64)
flow_size = flows.size()[2:]
vis_data = []
for pred_flow in net_outs:
vis_data.append(pred_flow)
deform_gt = deform_op(imgs_2, flows)
flows_pred = [
F.interpolate(x,
size=flow_size,
mode='bilinear',
align_corners=False)
for x in net_outs
]
deform_preds = [
deform_op(imgs_2, x) for x in flows_pred
]
vis_data.append(
F.interpolate(flows, size=tgt_size, mode='area'))
vis_data = [
F.interpolate(flow[:1], size=tgt_size)
for flow in vis_data
]
vis_data = [
fz.convert_from_flow(
flow[0].data.cpu().numpy().transpose(
1, 2, 0)).transpose(
2, 0, 1).astype('float32') / 255
for flow in vis_data
]
def convert_image(image):
image = F.interpolate(image,
size=tgt_size,
mode='area')
image = image[0]
image = image.data.cpu().numpy()
image = image[::-1]
image = image.transpose(1, 2, 0)
image = image * np.array(STD) + np.array(MEANS)
image = image.transpose(2, 0, 1)
image = image / 255
image = np.clip(image, -1, 1)
image = image[::-1]
return image
vis_data.append(convert_image(imgs_1))
vis_data.append(convert_image(imgs_2))
vis_data.append(convert_image(deform_gt))
vis_data.extend(
[convert_image(x) for x in deform_preds])
vis_data_stack = np.stack(vis_data, axis=0)
w.add_images("preds_flow", vis_data_stack)
elif cfg.flow.warp_mode == "flow":
import flowiz as fz
tgt_size = (64, 64)
vis_data = []
for pred_flow, _, _ in net_outs["preds_flow"]:
vis_data.append(pred_flow)
vis_data = [
F.interpolate(flow[:1], size=tgt_size)
for flow in vis_data
]
vis_data = [
fz.convert_from_flow(
flow[0].data.cpu().numpy().transpose(
1, 2, 0)).transpose(
2, 0, 1).astype('float32') / 255
for flow in vis_data
]
input_image = F.interpolate(images,
size=tgt_size,
mode='area')
input_image = input_image[0]
input_image = input_image.data.cpu().numpy()
input_image = input_image.transpose(1, 2, 0)
input_image = input_image * np.array(
STD[::-1]) + np.array(MEANS[::-1])
input_image = input_image.transpose(2, 0, 1)
input_image = input_image / 255
input_image = np.clip(input_image, -1, 1)
vis_data.append(input_image)
vis_data_stack = np.stack(vis_data, axis=0)
w.add_images("preds_flow", vis_data_stack)
iteration += 1
w.set_step(iteration)
if rank == 0 and iteration % args.save_interval == 0 and iteration != args.start_iter:
if args.keep_latest:
latest = SavePath.get_latest(args.save_folder,
cfg.name)
logger.info('Saving state, iter: {}'.format(iteration))
yolact_net.save_weights(save_path(epoch, iteration))
if args.keep_latest and latest is not None:
if args.keep_latest_interval <= 0 or iteration % args.keep_latest_interval != args.save_interval:
logger.info('Deleting old save...')
os.remove(latest)
misc.barrier()
# This is done per epoch
if args.validation_epoch > 0:
if epoch % args.validation_epoch == 0 and epoch > 0:
if rank == 0:
compute_validation_map(yolact_net, val_dataset)
misc.barrier()
except KeyboardInterrupt:
misc.barrier()
if args.interrupt_no_save:
logger.info('No save on interrupt, just exiting...')
elif rank == 0:
print('Stopping early. Saving network...')
# Delete previous copy of the interrupted network so we don't spam the weights folder
SavePath.remove_interrupt(args.save_folder)
yolact_net.save_weights(
save_path(epoch,
repr(iteration) + '_interrupt'))
return
if rank == 0:
yolact_net.save_weights(save_path(epoch, iteration))
def main(argv):
# set fixed random seed, load config files
tf.random.set_seed(RANDOM_SEED)
# using mix precision or not
if MIXPRECISION:
policy = mixed_precision.Policy('mixed_float16')
mixed_precision.set_policy(policy)
# get params for model
train_iter, input_size, num_cls, lrs_schedule_params, loss_params, parser_params, model_params = get_params(
FLAGS.name)
# -----------------------------------------------------------------
# set up Grappler for graph optimization
# Ref: https://www.tensorflow.org/guide/graph_optimization
@contextlib.contextmanager
def options(opts):
old_opts = tf.config.optimizer.get_experimental_options()
tf.config.optimizer.set_experimental_options(opts)
try:
yield
finally:
tf.config.optimizer.set_experimental_options(old_opts)
# -----------------------------------------------------------------
# Creating the instance of the model specified.
logging.info("Creating the model instance of YOLACT")
model = Yolact(**model_params)
# add weight decay
for layer in model.layers:
if isinstance(layer, tf.keras.layers.Conv2D) or isinstance(
layer, tf.keras.layers.Dense):
layer.add_loss(lambda: tf.keras.regularizers.l2(FLAGS.weight_decay)
(layer.kernel))
if hasattr(layer, 'bias_regularizer') and layer.use_bias:
layer.add_loss(lambda: tf.keras.regularizers.l2(FLAGS.weight_decay)
(layer.bias))
# -----------------------------------------------------------------
# Creating dataloaders for training and validation
logging.info("Creating the dataloader from: %s..." % FLAGS.tfrecord_dir)
dateset = ObjectDetectionDataset(dataset_name=FLAGS.name,
tfrecord_dir=os.path.join(
FLAGS.tfrecord_dir, FLAGS.name),
anchor_instance=model.anchor_instance,
**parser_params)
train_dataset = dateset.get_dataloader(subset='train',
batch_size=FLAGS.batch_size)
valid_dataset = dateset.get_dataloader(subset='val', batch_size=1)
# count number of valid data for progress bar
# Todo any better way to do it?
num_val = 0
for _ in valid_dataset:
num_val += 1
# -----------------------------------------------------------------
# Choose the Optimizor, Loss Function, and Metrics, learning rate schedule
lr_schedule = learning_rate_schedule.Yolact_LearningRateSchedule(
**lrs_schedule_params)
logging.info("Initiate the Optimizer and Loss function...")
optimizer = tf.keras.optimizers.SGD(learning_rate=lr_schedule,
momentum=FLAGS.momentum)
criterion = loss_yolact.YOLACTLoss(**loss_params)
train_loss = tf.keras.metrics.Mean('train_loss', dtype=tf.float32)
loc = tf.keras.metrics.Mean('loc_loss', dtype=tf.float32)
conf = tf.keras.metrics.Mean('conf_loss', dtype=tf.float32)
mask = tf.keras.metrics.Mean('mask_loss', dtype=tf.float32)
seg = tf.keras.metrics.Mean('seg_loss', dtype=tf.float32)
# -----------------------------------------------------------------
# Setup the TensorBoard for better visualization
# Ref: https://www.tensorflow.org/tensorboard/get_started
logging.info("Setup the TensorBoard...")
current_time = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
train_log_dir = './logs/gradient_tape/' + current_time + '/train'
test_log_dir = './logs/gradient_tape/' + current_time + '/test'
train_summary_writer = tf.summary.create_file_writer(train_log_dir)
test_summary_writer = tf.summary.create_file_writer(test_log_dir)
# -----------------------------------------------------------------
# Start the Training and Validation Process
logging.info("Start the training process...")
# setup checkpoints manager
checkpoint = tf.train.Checkpoint(step=tf.Variable(1),
optimizer=optimizer,
model=model)
manager = tf.train.CheckpointManager(checkpoint,
directory="./checkpoints",
max_to_keep=5)
# restore from latest checkpoint and iteration
status = checkpoint.restore(manager.latest_checkpoint)
if manager.latest_checkpoint:
logging.info("Restored from {}".format(manager.latest_checkpoint))
else:
logging.info("Initializing from scratch.")
best_masks_map = 0.
iterations = checkpoint.step.numpy()
for image, labels in train_dataset:
# check iteration and change the learning rate
if iterations > train_iter:
break
checkpoint.step.assign_add(1)
iterations += 1
with options({
'constant_folding': True,
'layout_optimize': True,
'loop_optimization': True,
'arithmetic_optimization': True,
'remapping': True
}):
loc_loss, conf_loss, mask_loss, seg_loss = train_step(
model, criterion, train_loss, optimizer, image, labels,
num_cls)
loc.update_state(loc_loss)
conf.update_state(conf_loss)
mask.update_state(mask_loss)
seg.update_state(seg_loss)
with train_summary_writer.as_default():
tf.summary.scalar('Total loss',
train_loss.result(),
step=iterations)
tf.summary.scalar('Loc loss', loc.result(), step=iterations)
tf.summary.scalar('Conf loss', conf.result(), step=iterations)
tf.summary.scalar('Mask loss', mask.result(), step=iterations)
tf.summary.scalar('Seg loss', seg.result(), step=iterations)
if iterations and iterations % FLAGS.print_interval == 0:
tf.print(
"Iteration {}, LR: {}, Total Loss: {}, B: {}, C: {}, M: {}, S:{} "
.format(iterations,
optimizer._decayed_lr(var_dtype=tf.float32),
train_loss.result(), loc.result(), conf.result(),
mask.result(), seg.result()))
if iterations and iterations % FLAGS.save_interval == 0:
# save checkpoint
save_path = manager.save()
logging.info("Saved checkpoint for step {}: {}".format(
int(checkpoint.step), save_path))
# validation and print mAP table
all_map = evaluate(model, valid_dataset, num_val, num_cls)
box_map, mask_map = all_map['box']['all'], all_map['mask']['all']
tf.print(f"box mAP:{box_map}, mask mAP:{mask_map}")
with test_summary_writer.as_default():
tf.summary.scalar('Box mAP', box_map, step=iterations)
tf.summary.scalar('Mask mAP', mask_map, step=iterations)
# Saving the weights:
if mask_map > best_masks_map:
best_masks_map = mask_map
model.save_weights(
f'{FLAGS.weights}/weights_{FLAGS.name}_{str(best_masks_map)}.h5'
)
# reset the metrics
train_loss.reset_states()
loc.reset_states()
conf.reset_states()
mask.reset_states()
seg.reset_states()
