'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)
# Define augmentation transformations as a composition
# composed_transforms = transforms.Compose([tr.RandomHorizontalFlip(),
# tr.ScaleNRotate(rots=(-30, 30), scales=(.75, 1.25)),
# tr.ToTensor()])
composed_transforms = transforms.Compose(
[tr.ToTensor(), tr.VideoResize([480, 848])])
#composed_transforms = transforms.Compose([tr.ToTensor()])
# composed_transforms = transforms.Compose([tr.ToTensor()])
# Testing dataset and its iterator
db_test = db.DAVIS2016(train=False,
train_online=False,
db_root_dir=db_root_dir,
transform=composed_transforms,
seq_name=seqname)
testloader = DataLoader(db_test,
batch_size=testBatch,
shuffle=False,
num_workers=2)
def getHeatMapFrom2DArray(inArray):
inArray_ = (inArray - np.min(inArray)) / (np.max(inArray) -
np.min(inArray))
cm = plt.get_cmap('jet')
retValue = cm(inArray_)
return retValue
{'params': [pr[1] for pr in net.score_dsn.named_parameters() if 'bias' in pr[0]], 'lr': 2 * lr / 10,
'initial_lr': 2 * lr / 10},
{'params': [pr[1] for pr in net.upscale.named_parameters() if 'weight' in pr[0]], 'lr': 0, 'initial_lr': 0},
{'params': [pr[1] for pr in net.upscale_.named_parameters() if 'weight' in pr[0]], 'lr': 0, 'initial_lr': 0},
{'params': net.fuse.weight, 'lr': lr / 100, 'initial_lr': lr / 100, 'weight_decay': wd},
{'params': net.fuse.bias, 'lr': 2 * lr / 100, 'initial_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, inputRes=None, db_root_dir=db_root_dir, transform=composed_transforms)
trainloader = DataLoader(db_train, batch_size=p['trainBatch'], shuffle=True, num_workers=2)
# Testing dataset and its iterator
db_test = db.DAVIS2016(train=False, db_root_dir=db_root_dir, transform=tr.ToTensor())
testloader = DataLoader(db_test, batch_size=testBatch, shuffle=False, num_workers=2)
num_img_tr = len(trainloader)
num_img_ts = len(testloader)
running_loss_tr = [0] * 5
running_loss_ts = [0] * 5
loss_tr = []
loss_ts = []
aveGrad = 0
print("Training Network")
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()