def train(train_path, model_save_path, num_epochs=3, model_load_path=None, patience=2, gamma=0.1, imp_thresh=0.001):
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
if model_load_path:
model = SSD_Flow()
state_dict = torch.load(model_load_path)
model.load_state_dict(state_dict['model_state_dict'])
else:
model = SSD_Flow(initialize=True)
model.freeze(freeze_add=False)
model.train_setbn(bn_eval=False) # need to set BN to eval mode if using batch size of 1
model = model.to(device)
trainable_params = [param for param in model.parameters() if param.requires_grad]
optimizer = torch.optim.SGD(trainable_params, lr=0.001, momentum=0.9, weight_decay=0.0005)
if model_load_path: optimizer.load_state_dict(state_dict['optimizer_state_dict'])
start_epoch = state_dict['epoch']+1 if model_load_path else 0
# prev_best_loss = state_dict['loss'] if model_load_path else float('inf')
# num_epochs_no_imp = 0
# imp_this_lr = False
dboxes = SSD_utils.dboxes300_coco()
loss_criterion = Loss(dboxes)
scene_files = sorted(os.listdir(train_path))
for epoch in range(start_epoch, num_epochs):
epoch_loss = 0
for scene_num, record_file in enumerate(scene_files):
try:
scene_data = Dataset(os.path.join(train_path, record_file), dataset_pb2.CameraName.FRONT, dboxes, augment=False, flow=True)
except:
print(record_file, "Corrupted!")
continue
scene_dataloader = torch.utils.data.DataLoader(scene_data, batch_size=32, shuffle=False, collate_fn=collate_fn, drop_last=True)
scene_loss = 0
start = time()
for frame_num, (_, _, images, bboxes, labels, flows) in enumerate(scene_dataloader):
if images is None or len(labels) == 1: continue # batch size of 1 will cause error
images = images.to(device)
bboxes = bboxes.to(device)
labels = labels.to(device)
flows = flows.to(device)
optimizer.zero_grad()
pred_boxes, pred_scores = model(images, flows)
loss = loss_criterion(pred_boxes, pred_scores, bboxes.transpose(1,2), labels)
loss.backward()
optimizer.step()
scene_loss += loss.item()
scene_loss /= len(scene_data)
epoch_loss += scene_loss
print("Scene", scene_num, "loss:", scene_loss)
if (scene_num+1)%25==0:
torch.save({'epoch': epoch, 'scene_num': scene_num, 'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(), 'loss': scene_loss},
os.path.join(model_save_path, "ep_"+str(epoch)+'_sc_'+str(scene_num)))
epoch_loss /= len(scene_files)
print("\nEpoch", epoch, "loss:", epoch_loss, "\n")
torch.save({'epoch': epoch, 'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(), 'loss': epoch_loss},
os.path.join(model_save_path, "ep_"+str(epoch)))
np.random.shuffle(scene_files)
else:
im = self.transform(im)
bboxes, labels = self.encoder.encode(bboxes, labels)
if self.flows is not None:
return orig_size, index, im, bboxes, labels, self.flows[
index]
else:
return orig_size, index, im, bboxes, labels
# necessary to discard samples without any objects
def collate_fn(batch):
batch = list(filter(lambda data: data is not None, batch))
if len(batch) > 0:
return torch.utils.data.dataloader.default_collate(batch)
else:
return None, None, None, None, None
if __name__ == '__main__':
tf.enable_eager_execution()
filename = 'data/train/segment-10206293520369375008_2796_800_2816_800_with_camera_labels.tfrecord'
dboxes = SSD_utils.dboxes300_coco()
dataset = Dataset(filename, open_dataset.CameraName.FRONT, dboxes)
for i in range(190):
if i != 150: continue
labels = dataset.__getitem__(i)[2]
#print(torch.sum(labels), labels.shape)
def inference(model_path,
eval_path,
pred_save_path,
view=dataset_pb2.CameraName.FRONT):
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
with torch.no_grad():
model = SSD_Ablation()
state_dict = torch.load(model_path)
model.load_state_dict(state_dict['model_state_dict'])
model.eval()
model = model.to(device)
dboxes = SSD_utils.dboxes300_coco()
decoder = SSD_utils.Encoder(dboxes)
scene_files = sorted(os.listdir(eval_path))
pred_objects = metrics_pb2.Objects()
for scene_num, record_file in enumerate(scene_files):
scene_data = Dataset(os.path.join(eval_path, record_file),
view,
dboxes,
val=True,
flow=False)
scene_dataloader = torch.utils.data.DataLoader(
scene_data,
batch_size=32,
shuffle=False,
collate_fn=collate_fn)
for batch_num, (orig_size, idxs, images, bboxes,
labels) in enumerate(scene_dataloader):
if images is None:
print("None")
continue
images = images.to(device)
bboxes = bboxes.to(device)
labels = labels.to(device)
#flows = flows.to(device)
pred_boxes, pred_scores = model(images)
preds = decoder.decode_batch(pred_boxes,
pred_scores,
criteria=0.5,
max_output=100)
if metrics_type == 'lib':
for i, pred in enumerate(preds):
create_prediction_file(pred, orig_size[i], record_file,
idxs[i], pred_save_path)
elif metrics_type == 'native':
for i, pred in enumerate(preds):
frame = scene_data.get_frame(idxs[i])
for j in range(len(pred[0])):
pred_objects.objects.append(
create_prediction_obj(frame, pred[0][j],
pred[2][j], pred[1][j],
view, orig_size[i]))
f = open(pred_save_path, 'wb')
f.write(pred_objects.SerializeToString())
f.close()
print("Finished scene", scene_num)
def train(train_path,
model_save_path,
num_epochs=3,
model_load_path=None,
patience=2,
gamma=0.1,
imp_thresh=0.001,
set_size=25):
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
if model_load_path:
model = SSD_Baseline()
state_dict = torch.load(model_load_path)
model.load_state_dict(state_dict['model_state_dict'])
else:
model = SSD_Baseline(initialize=True)
model.freeze()
model.train_setbn(
bn_eval=False) # need to set BN to eval mode if using batch size of 1
model = model.to(device)
trainable_params = [
param for param in model.parameters() if param.requires_grad
]
optimizer = torch.optim.SGD(trainable_params,
lr=0.001,
momentum=0.9,
weight_decay=0.0005)
if model_load_path:
optimizer.load_state_dict(state_dict['optimizer_state_dict'])
for g in optimizer.param_groups:
print(g['lr'])
sys.exit()
start_epoch = state_dict['epoch'] + 1 if model_load_path else 0
prev_best_loss = state_dict['loss'] if model_load_path else float('inf')
num_sets_no_imp = 0
imp_this_lr = False
dboxes = SSD_utils.dboxes300_coco()
loss_criterion = Loss(dboxes)
scene_files = sorted(os.listdir(train_path))
for epoch in range(start_epoch, num_epochs):
np.random.shuffle(scene_files)
epoch_loss = 0
set_loss = 0
for scene_num, record_file in enumerate(scene_files):
print(record_file)
try:
scene_data = Dataset(os.path.join(train_path, record_file),
dataset_pb2.CameraName.FRONT,
dboxes,
augment=False)
except:
print("Corrupted!")
continue
scene_dataloader = torch.utils.data.DataLoader(
scene_data,
batch_size=32,
shuffle=True,
collate_fn=collate_fn,
drop_last=True)
scene_loss = 0
start = time()
for frame_num, (_, _, images, bboxes,
labels) in enumerate(scene_dataloader):
if len(labels) == 1:
continue # batch size of 1 will cause error
images = images.to(device)
bboxes = bboxes.to(device)
labels = labels.to(device)
optimizer.zero_grad()
pred_boxes, pred_scores = model(images)
loss = loss_criterion(pred_boxes, pred_scores,
bboxes.transpose(1, 2), labels)
loss.backward()
optimizer.step()
scene_loss += loss.item()
scene_loss /= len(scene_data)
epoch_loss += scene_loss
set_loss += scene_loss
print("Scene", scene_num, "loss:", scene_loss)
if (scene_num + 1) % set_size == 0:
set_loss /= set_size
print("\nSet loss", set_loss, '\n')
torch.save(
{
'epoch': epoch,
'set': scene_num // set_size,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'loss': set_loss
},
os.path.join(model_save_path,
"set_" + str(scene_num // set_size)))
# LR stepping and early stopping
if set_loss > prev_best_loss - imp_thresh:
num_sets_no_imp += 1
if num_sets_no_imp == patience:
if imp_this_lr:
for g in optimizer.param_groups:
g['lr'] *= gamma
imp_this_lr = False
else:
return
else:
num_sets_no_imp = 0
imp_this_lr = True
prev_best_loss = set_loss
set_loss = 0
epoch_loss /= len(scene_files)
print("\nEpoch", epoch, "loss:", epoch_loss, "\n")
torch.save(
{
'epoch': epoch,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'loss': epoch_loss
}, os.path.join(model_save_path, "epoch_" + str(epoch)))