import torch
import torch.nn as nn
from base_learner import BaseLearner
from model.pytorch_lstm import *
from torch.utils.data import DataLoader
from smoke_video_dataset import SmokeVideoDataset
import torch.optim as optim
from torch.optim import lr_scheduler
import torch.nn.functional as F
from torch.autograd import Variable
import uuid
from sklearn.metrics import classification_report
from sklearn.metrics import accuracy_score
import numpy as np
import random
from torchvision.transforms import Compose
from video_transforms import RandomResizedCrop, RandomHorizontalFlip, ColorJitter, RandomPerspective, RandomErasing, Resize, Normalize
from torch.utils.tensorboard import SummaryWriter
import json
# Conv LSTM learner
class LSTMLearner(BaseLearner):
def __init__(self,
batch_size=8, # Batch 8 per GPU
lr=0.01,
max_steps=64e3,
momentum=0.9,
#milestones = [40, 500, 1000],
gamma = 0.1,
weight_decay = 0.000001,
num_workers=1,
num_of_action_classes=2,
num_steps_per_update=2,
num_steps_per_check=10,
#use_cuda=torch.cuda.is_available(),
use_cuda=True,
parallel=True,
augment=True,
save_model_path="../data/saved_lstm/", # path for saving the models
train_writer_p="../data/ts_runs/lstm_train",
val_writer_p="../data/ts_runs/lstm_val"
):
super().__init__()
self.create_logger(log_path="LSTMLearner_test.log")
self.log("Use LSTM learner")
self.batch_size = batch_size
self.lr = lr
self.max_steps = max_steps
self.momentum = momentum
#self.milestones = milestones
self.gamma = gamma
self.weight_decay = weight_decay
self.num_workers = num_workers
self.num_of_action_classes = num_of_action_classes
self.num_steps_per_update = num_steps_per_update
self.num_steps_per_check = num_steps_per_check
self.use_cuda = use_cuda
self.parallel = parallel
self.augment = augment
self.save_model_path = save_model_path
self.train_writer_p = train_writer_p
self.val_writer_p = val_writer_p
def random_frames_from_batch(self, data):
b, c, f, h, w = list(data.shape)
new_batch = torch.zeros(b, c, h, w)
for i in range(b):
frame = random.randint(0, f - 1)
new_batch[i, :, :, :] = data[i, :, frame, :, :]
return new_batch
def compress_videos_to_frames(self, n):
b, c, t, h, w = n.shape
ret = torch.zeros(b, c*t, h, w)
frame = 0
for channel in range(c*t):
ret[:,channel,:,:] = n[:,channel%c, frame,:,:]
if channel > 0 and channel % c == 0: frame += 1
return ret
def set_dataloader(self, metadata_path, p_vid, mode, tf):
dataloader = {}
for phase in metadata_path:
self.log("Create dataloader for " + phase)
dataset = SmokeVideoDataset(metadata_path=metadata_path[phase], root_dir=p_vid, transform=tf[phase]) #mode=mode, transform=tf[phase])
dataloader[phase] = DataLoader(dataset, batch_size=self.batch_size, shuffle=True, num_workers=self.num_workers, pin_memory=True)
return dataloader
def labels_to_list(self, labels):
# Convert labels obtained from the dataloader to a list of action classes
return np.argmax(labels.numpy().max(axis=2), axis=1).tolist()
def to_variable(self, v):
if self.use_cuda:
v = v.cuda() # move to gpu
return Variable(v)
def make_pred(self, model, frames):
# Upsample prediction to frame length (because we want prediction for each frame)
return F.interpolate(model(frames), frames.size(2), mode="linear", align_corners=True)
def fit(self,
mode="rgb",
p_metadata_train="../data/split/metadata_train_split_0_by_camera.json",
p_metadata_validation="../data/split/metadata_validation_split_0_by_camera.json",
p_vid="../data/"):
self.log("="*60)
self.log("="*60)
self.log("Start training...")
if mode == "rgb":
model = CNN_LSTM()
p_vid = p_vid + "rgb/"
elif mode == "flow":
model = CNN_LSTM()
p_vid = p_vid + "flow/"
#self.load(model, "../data/saved_lstm/b1ca514-lstm-rgb/12912.pt")
if self.use_cuda:
model.cuda()
if self.parallel and torch.cuda.device_count() > 1:
self.log("Let's use " + str(torch.cuda.device_count()) + " GPUs!")
model = nn.DataParallel(model)
writer_train = SummaryWriter(self.train_writer_p)
writer_val = SummaryWriter(self.val_writer_p)
# Load datasets
metadata_path = {"train": p_metadata_train, "validation": p_metadata_validation}
ts = self.get_transform()
transform = {"train": ts, "validation": ts}
if self.augment:
transform["train"] = self.get_transform(phase="train")
dataloader = self.set_dataloader(metadata_path, p_vid, mode, transform)
# Set optimizer
optimizer = optim.SGD(params=model.parameters(), lr=self.lr, momentum=self.momentum, weight_decay = self.weight_decay)
# lr_sche = optim.lr_scheduler.MultiStepLR(optimizer, milestones=self.milestones, gamma=self.gamma)
# Set Loss Fum nction
criterion = nn.BCEWithLogitsLoss()
# Set logging format
log_fm = "%s step: %d lr: %r loss: %.4f"
# Train and validate
steps = 0
epochs = 0
nspu = self.num_steps_per_update
nspc = self.num_steps_per_check
nspu_nspc = nspu * nspc
model_id = str(uuid.uuid4())[0:7] + "-lstm-" + mode
accum = {} # counter for accumulating gradients
tot_loss = {} # total loss
tot_loc_loss = {} # total localization loss
tot_cls_loss = {} # total classification loss
pred_labels = {} # predicted labels
true_labels = {} # true labels
for phase in ["train", "validation"]:
accum[phase] = 0
tot_loss[phase] = 0.0
tot_loc_loss[phase] = 0.0
tot_cls_loss[phase] = 0.0
pred_labels[phase] = []
true_labels[phase] = []
while steps < self.max_steps:
self.log("-"*40)
# Each epoch has a training and validation phase
for phase in ["train", "validation"]:
self.log("phase " + phase)
if phase == "train":
epochs += 1
self.log("epochs: %d steps: %d/%d" % (epochs, steps, self.max_steps))
model.train(True) # set model to training mode
else:
model.train(False) # set model to evaluate mode
optimizer.zero_grad()
# Iterate over data
for d in dataloader[phase]:
accum[phase] += 1
# Get inputs
frames, labels = d["frames"], d["labels"]
#frames = self.compress_videos_to_frames(frames)
frames = self.to_variable(frames)
true_labels[phase] += self.labels_to_list(labels)
labels = self.to_variable(d["labels"])
pred = model(frames)
pred = pred.permute(0,2,1)
_, predicted = torch.max(pred, dim=1)
cpu_pred = predicted.cpu().detach().tolist()
for lst in cpu_pred:
pred_labels[phase].append(lst[-1])
# Compute classification loss (with max-pooling along time, batch x channel x time)
cls_loss = criterion(torch.max(pred, dim=2)[0], torch.max(labels, dim=2)[0])
tot_cls_loss[phase] += cls_loss.data
# Backprop
loss = cls_loss / nspu
tot_loss[phase] += loss.data
if phase == "train":
loss.backward()
# Accumulate gradients during training
if (accum[phase] == nspu) and phase == "train":
steps += 1
accum[phase] = 0
optimizer.step()
optimizer.zero_grad()
# lr_sche.step()
if steps % nspc == 0:
#lr = lr_sche.get_lr()[0]
lr = self.lr
tl = tot_loss[phase]/nspc
writer_train.add_scalar("Training Loss", tl, global_step=steps)
self.log(log_fm % (phase, steps, lr, tl))
tot_loss[phase] = 0.0
if phase == "validation":
# lr = lr_sche.get_lr()[0]
lr = self.lr
tl = (tot_loss[phase]*nspu)/accum[phase]
writer_val.add_scalar("Validation Loss", tl, global_step=steps)
self.log(log_fm % (phase, steps, lr, tl))
tot_loss[phase] = 0.0
accum[phase] = 0
p_model = self.save_model_path + model_id + "/" + str(steps) + ".pt"
self.save(model, p_model)
for phase in ["train", "validation"]:
self.log("Performance for " + phase)
accuracy = accuracy_score(true_labels[phase], pred_labels[phase])
if phase == "train": writer_train.add_scalar("Train Accuracy", accuracy, global_step=steps)
elif phase == "validation": writer_val.add_scalar("Val Accuracy", accuracy, global_step=steps)
self.log(classification_report(true_labels[phase], pred_labels[phase]))
pred_labels[phase] = []
true_labels[phase] = []
self.log("Done fit")
def test(self,
mode="rgb",
p_metadata_test="../data/split/metadata_test_split_0_by_camera.json",
p_vid="../data/",
p_model=None):
self.log("="*60)
self.log("="*60)
self.log("Start testing...")
if mode == "rgb":
model = CNN_LSTM()
p_vid = p_vid + "rgb/"
elif mode == "flow":
model = CNN_LSTM()
p_vid = p_vid + "flow/"
self.load(model, p_model)
if self.use_cuda:
model.cuda()
if self.parallel and torch.cuda.device_count() > 1:
self.log("Let's use " + str(torch.cuda.device_count()) + " GPUs!")
model = nn.DataParallel(model)
# Load dataset
metadata_path = {"test": p_metadata_test}
transform = {"test": self.get_transform()}
dataloader = self.set_dataloader(metadata_path, p_vid, mode, transform)
#self.load(model, p_model)
all_pred = {}
# Test
model.train(False)
true_labels = []
pred_labels = []
counter = 0
with torch.no_grad():
for d in dataloader["test"]:
if counter % 20 == 0:
self.log("Process batch " + str(counter))
counter += 1
# Get inputs
frames, labels = d["frames"], d["labels"]
names = d["file_name"]
#frames = self.compress_videos_to_frames(frames)
frames = self.to_variable(frames)
l = self.labels_to_list(labels)
true_labels += l
labels = self.to_variable(d["labels"])
pred = model(frames)
pred = pred.permute(0,2,1)
all_pred[counter] = {}
p = pred.cpu().detach()
_, predicted = torch.max(pred, dim=1)
cpu_pred = predicted.cpu().detach().tolist()
for i, lst in enumerate(cpu_pred):
all_pred[counter][i] = []
all_pred[counter][i].append("neg")
all_pred[counter][i].append(p[i,0,-1].item())
all_pred[counter][i].append("pos")
all_pred[counter][i].append(p[i,1,-1].item())
all_pred[counter][i].append("actual")
all_pred[counter][i].append(l[i])
all_pred[counter][i].append("name")
all_pred[counter][i].append(names[i])
pred_labels.append(lst[-1])
#frames = self.to_variable(frames)
#true_labels += self.labels_to_list(labels)
#labels = self.to_variable(d["labels"])
#pred = model(frames)
#pred.permute(0,2,1)
#_, predicted = torch.max(pred, dim=1)
#cpu_pred = predicted.cpu().detach().tolist()
#for lst in cpu_pred:
# pred_labels.append(lst[-1])
self.log(classification_report(true_labels, pred_labels))
with open ('lstm_pred.json', 'w') as f:
json.dump(all_pred, f)
self.log("Done test")
