# One training epoch
running_loss_tr = 0
np.random.seed(seed + epoch)
for ii, sample_batched in enumerate(trainloader):
inputs, gts = sample_batched['image'], sample_batched['gt']
# Forward-Backward of the mini-batch
inputs.requires_grad_()
inputs, gts = inputs.to(device), gts.to(device)
outputs = net.forward(inputs)
# Compute the fuse loss
loss = class_balanced_cross_entropy_loss(outputs[-1],
gts,
size_average=False)
running_loss_tr += loss.item() # PyTorch 0.4.0 style
# Print stuff
if epoch % (nEpochs // 20) == (nEpochs // 20 - 1):
running_loss_tr /= num_img_tr
loss_tr.append(running_loss_tr)
print('[Epoch: %d, numImages: %5d]' % (epoch + 1, ii + 1))
print('Loss: %f' % running_loss_tr)
writer.add_scalar('data/total_loss_epoch', running_loss_tr, epoch)
# Backward the averaged gradient
loss /= nAveGrad
loss.backward()
start_time = timeit.default_timer()
# One training epoch
for ii, sample_batched in enumerate(trainloader):
inputs, gts = sample_batched['image'], sample_batched['gt']
# Forward-Backward of the mini-batch
inputs.requires_grad_()
inputs, gts = inputs.to(device), gts.to(device)
outputs = net.forward(inputs)
# Compute the losses, side outputs and fuse
losses = [0] * len(outputs)
for i in range(0, len(outputs)):
losses[i] = class_balanced_cross_entropy_loss(outputs[i], gts, size_average=False)
running_loss_tr[i] += losses[i].item()
loss = (1 - epoch / nEpochs)*sum(losses[:-1]) + losses[-1]
# Print stuff
if ii % num_img_tr == num_img_tr - 1:
running_loss_tr = [x / num_img_tr for x in running_loss_tr]
loss_tr.append(running_loss_tr[-1])
writer.add_scalar('data/total_loss_epoch', running_loss_tr[-1], epoch)
print('[Epoch: %d, numImages: %5d]' % (epoch, ii + 1))
for l in range(0, len(running_loss_tr)):
print('Loss %d: %f' % (l, running_loss_tr[l]))
running_loss_tr[l] = 0
stop_time = timeit.default_timer()
print("Execution time: " + str(stop_time - start_time))
def train(epochs_wo_avegrad):
# Setting of parameters
if 'SEQ_NAME' not in os.environ.keys():
seq_name = 'blackswan'
else:
seq_name = str(os.environ['SEQ_NAME'])
db_root_dir = Path.db_root_dir()
save_dir = Path.save_root_dir()
if not os.path.exists(save_dir):
os.makedirs(os.path.join(save_dir))
vis_net = 0 # Visualize the network?
vis_res = 0 # Visualize the results?
nAveGrad = 5 # Average the gradient every nAveGrad iterations
nEpochs = epochs_wo_avegrad * nAveGrad # Number of epochs for training #CHANGED from 2000
snapshot = nEpochs # Store a model every snapshot epochs
parentEpoch = 240
# Parameters in p are used for the name of the model
p = {
'trainBatch': 1, # Number of Images in each mini-batch
}
seed = 0
parentModelName = 'parent'
# Select which GPU, -1 if CPU
gpu_id = 0
device = torch.device("cuda:" +
str(gpu_id) if torch.cuda.is_available() else "cpu")
# Network definition
net = vo.OSVOS(pretrained=0)
net.load_state_dict(
torch.load(os.path.join(
save_dir,
parentModelName + '_epoch-' + str(parentEpoch - 1) + '.pth'),
map_location=lambda storage, loc: storage))
# Logging into Tensorboard
log_dir = os.path.join(
save_dir, 'runs',
datetime.now().strftime('%b%d_%H-%M-%S') + '_' + socket.gethostname() +
'-' + seq_name)
writer = SummaryWriter(logdir=log_dir)
net.to(device) # PyTorch 0.4.0 style
# Visualize the network
if vis_net:
x = torch.randn(1, 3, 480, 854)
x.requires_grad_()
x = x.to(device)
y = net.forward(x)
g = viz.make_dot(y, net.state_dict())
g.view()
# Use the following optimizer
lr = 1e-8
wd = 0.0002
optimizer = optim.SGD([
{
'params': [
pr[1]
for pr in net.stages.named_parameters() if 'weight' in pr[0]
],
'weight_decay':
wd
},
{
'params':
[pr[1] for pr in net.stages.named_parameters() if 'bias' in pr[0]],
'lr':
lr * 2
},
{
'params': [
pr[1]
for pr in net.side_prep.named_parameters() if 'weight' in pr[0]
],
'weight_decay':
wd
},
{
'params': [
pr[1]
for pr in net.side_prep.named_parameters() if 'bias' in pr[0]
],
'lr':
lr * 2
},
{
'params': [
pr[1]
for pr in net.upscale.named_parameters() if 'weight' in pr[0]
],
'lr':
0
},
{
'params': [
pr[1]
for pr in net.upscale_.named_parameters() if 'weight' in pr[0]
],
'lr':
0
},
{
'params': net.fuse.weight,
'lr': lr / 100,
'weight_decay': wd
},
{
'params': net.fuse.bias,
'lr': 2 * lr / 100
},
],
lr=lr,
momentum=0.9)
# Preparation of the data loaders
# Define augmentation transformations as a composition
composed_transforms = transforms.Compose([
tr.RandomHorizontalFlip(),
tr.ScaleNRotate(rots=(-30, 30), scales=(.75, 1.25)),
tr.ToTensor()
])
# Training dataset and its iterator
db_train = db.DAVIS2016(train=True,
db_root_dir=db_root_dir,
transform=composed_transforms,
seq_name=seq_name)
trainloader = DataLoader(db_train,
batch_size=p['trainBatch'],
shuffle=True,
num_workers=1)
# Testing dataset and its iterator
db_test = db.DAVIS2016(train=False,
db_root_dir=db_root_dir,
transform=tr.ToTensor(),
seq_name=seq_name)
testloader = DataLoader(db_test,
batch_size=1,
shuffle=False,
num_workers=1)
num_img_tr = len(trainloader)
num_img_ts = len(testloader)
loss_tr = []
aveGrad = 0
print("Start of Online Training, sequence: " + seq_name)
start_time = timeit.default_timer()
# Main Training and Testing Loop
for epoch in range(0, nEpochs):
# One training epoch
running_loss_tr = 0
np.random.seed(seed + epoch)
for ii, sample_batched in enumerate(trainloader):
inputs, gts = sample_batched['image'], sample_batched['gt']
# Forward-Backward of the mini-batch
inputs.requires_grad_()
inputs, gts = inputs.to(device), gts.to(device)
outputs = net.forward(inputs)
# Compute the fuse loss
loss = class_balanced_cross_entropy_loss(outputs[-1],
gts,
size_average=False)
running_loss_tr += loss.item() # PyTorch 0.4.0 style
# Print stuff
if epoch % (nEpochs // 20) == (nEpochs // 20 - 1):
running_loss_tr /= num_img_tr
loss_tr.append(running_loss_tr)
print('[Epoch: %d, numImages: %5d]' % (epoch + 1, ii + 1))
print('Loss: %f' % running_loss_tr)
writer.add_scalar('data/total_loss_epoch', running_loss_tr,
epoch)
# Backward the averaged gradient
loss /= nAveGrad
loss.backward()
aveGrad += 1
# Update the weights once in nAveGrad forward passes
if aveGrad % nAveGrad == 0:
writer.add_scalar('data/total_loss_iter', loss.item(),
ii + num_img_tr * epoch)
optimizer.step()
optimizer.zero_grad()
aveGrad = 0
# Save the model
if (epoch % snapshot) == snapshot - 1 and epoch != 0:
torch.save(
net.state_dict(),
os.path.join(save_dir,
seq_name + '_epoch-' + str(epoch) + '.pth'))
stop_time = timeit.default_timer()
print('Online training time: ' + str(stop_time - start_time))
# Testing Phase
if vis_res:
import matplotlib.pyplot as plt
plt.close("all")
plt.ion()
f, ax_arr = plt.subplots(1, 3)
save_dir_res = os.path.join(save_dir, 'Results', seq_name)
if not os.path.exists(save_dir_res):
os.makedirs(save_dir_res)
print('Testing Network')
with torch.no_grad(): # PyTorch 0.4.0 style
# Main Testing Loop
for ii, sample_batched in enumerate(testloader):
img, gt, fname = sample_batched['image'], sample_batched[
'gt'], sample_batched['fname']
# Forward of the mini-batch
inputs, gts = img.to(device), gt.to(device)
outputs = net.forward(inputs)
for jj in range(int(inputs.size()[0])):
pred = np.transpose(
outputs[-1].cpu().data.numpy()[jj, :, :, :], (1, 2, 0))
pred = 1 / (1 + np.exp(-pred))
pred = np.squeeze(pred)
# Save the result, attention to the index jj
sm.imsave(
os.path.join(save_dir_res,
os.path.basename(fname[jj]) + '.png'), pred)
if vis_res:
img_ = np.transpose(img.numpy()[jj, :, :, :], (1, 2, 0))
gt_ = np.transpose(gt.numpy()[jj, :, :, :], (1, 2, 0))
gt_ = np.squeeze(gt)
# Plot the particular example
ax_arr[0].cla()
ax_arr[1].cla()
ax_arr[2].cla()
ax_arr[0].set_title('Input Image')
ax_arr[1].set_title('Ground Truth')
ax_arr[2].set_title('Detection')
ax_arr[0].imshow(im_normalize(img_))
ax_arr[1].imshow(gt_)
ax_arr[2].imshow(im_normalize(pred))
plt.pause(0.001)
writer.close()
def train(self,
first_frame,
n_interaction,
obj_id,
scribbles_data,
scribble_iter,
subset,
use_previous_mask=False):
nAveGrad = 1
num_workers = 4
train_batch = min(n_interaction, self.train_batch)
frames_list = interactive_utils.scribbles.annotated_frames_object(
scribbles_data, obj_id)
scribbles_list = scribbles_data['scribbles']
seq_name = scribbles_data['sequence']
if obj_id == 1 and n_interaction == 1:
self.prev_models = {}
# Network definition
if n_interaction == 1:
print('Loading weights from: {}'.format(self.parent_model))
self.net.load_state_dict(self.parent_model_state)
self.prev_models[obj_id] = None
else:
print(
'Loading weights from previous network: objId-{}_interaction-{}_scribble-{}.pth'
.format(obj_id, n_interaction - 1, scribble_iter))
self.net.load_state_dict(self.prev_models[obj_id])
lr = 1e-8
wd = 0.0002
optimizer = optim.SGD([
{
'params': [
pr[1] for pr in self.net.stages.named_parameters()
if 'weight' in pr[0]
],
'weight_decay':
wd
},
{
'params': [
pr[1] for pr in self.net.stages.named_parameters()
if 'bias' in pr[0]
],
'lr':
lr * 2
},
{
'params': [
pr[1] for pr in self.net.side_prep.named_parameters()
if 'weight' in pr[0]
],
'weight_decay':
wd
},
{
'params': [
pr[1] for pr in self.net.side_prep.named_parameters()
if 'bias' in pr[0]
],
'lr':
lr * 2
},
{
'params': [
pr[1] for pr in self.net.upscale.named_parameters()
if 'weight' in pr[0]
],
'lr':
0
},
{
'params': [
pr[1] for pr in self.net.upscale_.named_parameters()
if 'weight' in pr[0]
],
'lr':
0
},
{
'params': self.net.fuse.weight,
'lr': lr / 100,
'weight_decay': wd
},
{
'params': self.net.fuse.bias,
'lr': 2 * lr / 100
},
],
lr=lr,
momentum=0.9)
prev_mask_path = os.path.join(
self.save_res_dir, 'interaction-{}'.format(n_interaction - 1),
'scribble-{}'.format(scribble_iter))
composed_transforms_tr = transforms.Compose([
tr.SubtractMeanImage(self.meanval),
tr.CustomScribbleInteractive(scribbles_list,
first_frame,
use_previous_mask=use_previous_mask,
previous_mask_path=prev_mask_path),
tr.RandomHorizontalFlip(),
tr.ScaleNRotate(rots=(-30, 30), scales=(.75, 1.25)),
tr.ToTensor()
])
# Training dataset and its iterator
db_train = db.DAVIS2017(split=subset,
transform=composed_transforms_tr,
custom_frames=frames_list,
seq_name=seq_name,
obj_id=obj_id,
no_gt=True,
retname=True)
trainloader = DataLoader(db_train,
batch_size=train_batch,
shuffle=True,
num_workers=num_workers)
num_img_tr = len(trainloader)
loss_tr = []
aveGrad = 0
start_time = timeit.default_timer()
# Main Training and Testing Loop
epoch = 0
while 1:
# One training epoch
running_loss_tr = 0
for ii, sample_batched in enumerate(trainloader):
inputs, gts, void = sample_batched['image'], sample_batched[
'scribble_gt'], sample_batched['scribble_void_pixels']
# Forward-Backward of the mini-batch
inputs, gts, void = Variable(inputs), Variable(gts), Variable(
void)
if self.gpu_id >= 0:
inputs, gts, void = inputs.cuda(), gts.cuda(), void.cuda()
outputs = self.net.forward(inputs)
# Compute the fuse loss
loss = class_balanced_cross_entropy_loss(outputs[-1],
gts,
size_average=False,
void_pixels=void)
running_loss_tr += loss.item()
# Print stuff
if epoch % 10 == 0:
running_loss_tr /= num_img_tr
loss_tr.append(running_loss_tr)
print('[Epoch: %d, numImages: %5d]' % (epoch + 1, ii + 1))
print('Loss: %f' % running_loss_tr)
# writer.add_scalar('data/total_loss_epoch', running_loss_tr, epoch)
# Backward the averaged gradient
loss /= nAveGrad
loss.backward()
aveGrad += 1
# Update the weights once in nAveGrad forward passes
if aveGrad % nAveGrad == 0:
# writer.add_scalar('data/total_loss_iter', loss.data[0], ii + num_img_tr * epoch)
optimizer.step()
optimizer.zero_grad()
aveGrad = 0
epoch += train_batch
stop_time = timeit.default_timer()
if stop_time - start_time > self.time_budget:
break
# Save the model into dictionary
self.prev_models[obj_id] = copy.deepcopy(self.net.state_dict())
def train(self, first_frame, n_interaction, obj_id, scribbles_data, scribble_iter, subset, use_previous_mask=False):
nAveGrad = 1
num_workers = 4
train_batch = min(n_interaction, self.train_batch)
frames_list = interactive_utils.scribbles.annotated_frames_object(scribbles_data, obj_id)
scribbles_list = scribbles_data['scribbles']
seq_name = scribbles_data['sequence']
if obj_id == 1 and n_interaction == 1:
self.prev_models = {}
# # Network definition
# if n_interaction == 1:
# print('Loading weights from: {}'.format(self.parent_model))
# self.net.load_state_dict(self.parent_model_state)
# self.prev_models[obj_id] = None
# else:
# print('Loading weights from previous network: objId-{}_interaction-{}_scribble-{}.pth'
# .format(obj_id, n_interaction-1, scribble_iter))
# self.net.load_state_dict(self.prev_models[obj_id])
lr = 1e-5
wd = 0.0002
# optimizer = optim.SGD([
# {'params': [pr[1] for pr in self.net.stages.named_parameters() if 'weight' in pr[0]], 'weight_decay': wd},
# {'params': [pr[1] for pr in self.net.stages.named_parameters() if 'bias' in pr[0]], 'lr': lr * 2},
# {'params': [pr[1] for pr in self.net.side_prep.named_parameters() if 'weight' in pr[0]], 'weight_decay': wd},
# {'params': [pr[1] for pr in self.net.side_prep.named_parameters() if 'bias' in pr[0]], 'lr': lr * 2},
# {'params': [pr[1] for pr in self.net.upscale.named_parameters() if 'weight' in pr[0]], 'lr': 0},
# {'params': [pr[1] for pr in self.net.upscale_.named_parameters() if 'weight' in pr[0]], 'lr': 0},
# {'params': self.net.fuse.weight, 'lr': lr / 100, 'weight_decay': wd},
# {'params': self.net.fuse.bias, 'lr': 2 * lr / 100},
# ], lr=lr, momentum=0.9)
optimizer = optim.Adam(filter(lambda p: p.requires_grad, self.net.parameters()), lr=lr, momentum=0.9)
prev_mask_path = os.path.join(self.save_res_dir, 'interaction-{}'.format(n_interaction-1),
'scribble-{}'.format(scribble_iter))
composed_transforms_tr = transforms.Compose([tr.CenterCrop((480,832)),tr.SubtractMeanImage(self.meanval),
tr.CustomScribbleInteractive(scribbles_list, first_frame,
use_previous_mask=use_previous_mask,
previous_mask_path=prev_mask_path),
tr.RandomHorizontalFlip(),
tr.ScaleNRotate(rots=(-30, 30), scales=(.75, 1.25)),
tr.ToTensor()])
# Training dataset and its iterator
# db_train = db.DAVIS2017(split=subset, transform=composed_transforms_tr,
# custom_frames=frames_list, seq_name=seq_name,
# obj_id=obj_id, no_gt=True, retname=True)
db_train = db_scribblenet.DAVIS2017(split=subset, transform=composed_transforms_tr,
custom_frames=frames_list, seq_name=seq_name,
obj_id=obj_id, no_gt=True, retname=True)
trainloader = DataLoader(db_train, batch_size=train_batch, shuffle=True, num_workers=num_workers)
num_img_tr = len(trainloader)
loss_tr = []
aveGrad = 0
# List of all previous masks and aggregated features
prev_masks = []
prev_aggs = []
start_time = timeit.default_timer()
# Main Training and Testing Loop
epoch = 0
while 1:
# One training epoch
running_loss_tr = 0
for ii, sample_batched in enumerate(trainloader):
optimizer.zero_grad()
# Parse from dataset loader
inputs = sample_batched['images'].cuda()
gts = sample_batched['scribble_gt'].cuda()
scribbles = sample_batched['scribble_raw'].cuda()
scribbles_idx = sample_batched['scribble_idx'].cuda()
# Forward-Backward of the mini-batch
# prev_masks = torch.tensor(prev_masks).unsqueeze(0)
# prev_aggs = torch.tensor(prev_aggs).unsqueeze(0)
masks, agg = self.net.forward(inputs, scribbles, scribble_idx, prev_masks, prev_agg)
# Compute the fuse loss
loss = class_balanced_cross_entropy_loss(masks, gts, scribble_idx)
running_loss_tr += loss.item()
# Print stuff
if epoch % 10 == 0:
running_loss_tr /= num_img_tr
loss_tr.append(running_loss_tr)
print('[Epoch: %d, numImages: %5d]' % (epoch + 1, ii + 1))
print('Loss: %f' % running_loss_tr)
# Backward the averaged gradient
loss.backward()
optimizer.step()
# Update the current round data
prev_masks.append(masks.detach())
prev_aggs.append(agg.detach())
epoch += train_batch
stop_time = timeit.default_timer()
if stop_time - start_time > self.time_budget:
break
# Save the model into dictionary
self.prev_models[obj_id] = copy.deepcopy(self.net.state_dict())
