def train_model(dataset=dataset,
save_dir=save_dir,
num_classes=num_classes,
lr=lr,
num_epochs=nEpochs,
save_epoch=snapshot,
useTest=useTest,
test_interval=nTestInterval):
"""
Args:
num_classes (int): Number of classes in the data
num_epochs (int, optional): Number of epochs to train for.
"""
if modelName == 'C3D':
model = C3D_model.C3D(num_classes=num_classes, pretrained=True)
train_params = [{
'params': C3D_model.get_1x_lr_params(model),
'lr': lr
}, {
'params': C3D_model.get_10x_lr_params(model),
'lr': lr * 10
}]
elif modelName == 'R2Plus1D':
model = R2Plus1D_model.R2Plus1DClassifier(num_classes=num_classes,
layer_sizes=(2, 2, 2, 2))
train_params = [{
'params': R2Plus1D_model.get_1x_lr_params(model),
'lr': lr
}, {
'params': R2Plus1D_model.get_10x_lr_params(model),
'lr': lr * 10
}]
else:
print('We only implemented C3D and R2Plus1D models.')
raise NotImplementedError
criterion = nn.CrossEntropyLoss(
) # standard crossentropy loss for classification
optimizer = optim.SGD(train_params, lr=lr, momentum=0.9, weight_decay=5e-4)
scheduler = optim.lr_scheduler.StepLR(
optimizer, step_size=10,
gamma=0.1) # the scheduler divides the lr by 10 every 10 epochs
if resume_epoch == 0:
print("Training {} from scratch...".format(modelName))
else:
checkpoint = torch.load(os.path.join(
save_dir, 'models', modelName + '_' + dataset + '_epoch-' +
str(resume_epoch - 1) + '.pth.tar'),
map_location=lambda storage, loc: storage
) # Load all tensors onto the CPU
print("Initializing weights from: {}...".format(
os.path.join(
save_dir, 'models', modelName + '_' + dataset + '_epoch-' +
str(resume_epoch - 1) + '.pth.tar')))
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['opt_dict'])
print('Total params: %.2fM' %
(sum(p.numel() for p in model.parameters()) / 1000000.0))
model.to(device)
criterion.to(device)
log_dir = os.path.join(
save_dir, 'models',
datetime.now().strftime('%b%d_%H-%M-%S') + '_' + socket.gethostname())
writer = SummaryWriter(log_dir=log_dir)
print('Training model on {} dataset...'.format(dataset))
train_dataloader = DataLoader(VideoDataset(dataset=dataset,
split='train',
clip_len=16),
batch_size=16,
shuffle=True,
num_workers=4)
val_dataloader = DataLoader(VideoDataset(dataset=dataset,
split='val',
clip_len=16),
batch_size=16,
num_workers=4)
test_dataloader = DataLoader(VideoDataset(dataset=dataset,
split='test',
clip_len=16),
batch_size=16,
num_workers=4)
trainval_loaders = {'train': train_dataloader, 'val': val_dataloader}
trainval_sizes = {
x: len(trainval_loaders[x].dataset)
for x in ['train', 'val']
}
test_size = len(test_dataloader.dataset)
for epoch in range(resume_epoch, num_epochs):
# each epoch has a training and validation step
for phase in ['train', 'val']:
start_time = timeit.default_timer()
# reset the running loss and corrects
running_loss = 0.0
running_corrects = 0
# set model to train() or eval() mode depending on whether it is trained
# or being validated. Primarily affects layers such as BatchNorm or Dropout.
if phase == 'train':
# scheduler.step() is to be called once every epoch during training
scheduler.step()
model.train()
else:
model.eval()
for inputs, labels in tqdm(trainval_loaders[phase]):
# move inputs and labels to the device the training is taking place on
inputs = Variable(inputs, requires_grad=True).to(device)
labels = Variable(labels).to(device)
optimizer.zero_grad()
outputs = model(inputs)
probs = nn.Softmax(dim=1)(outputs)
preds = torch.max(probs, 1)[1]
loss = criterion(outputs, labels)
if phase == 'train':
loss.backward()
optimizer.step()
running_loss += loss.item() * inputs.size(0)
running_corrects += torch.sum(preds == labels.data)
epoch_loss = running_loss / trainval_sizes[phase]
epoch_acc = running_corrects.double() / trainval_sizes[phase]
if phase == 'train':
writer.add_scalar('data/train_loss_epoch', epoch_loss, epoch)
writer.add_scalar('data/train_acc_epoch', epoch_acc, epoch)
else:
writer.add_scalar('data/val_loss_epoch', epoch_loss, epoch)
writer.add_scalar('data/val_acc_epoch', epoch_acc, epoch)
print("[{}] Epoch: {}/{} Loss: {} Acc: {}".format(
phase, epoch + 1, nEpochs, epoch_loss, epoch_acc))
stop_time = timeit.default_timer()
print("Execution time: " + str(stop_time - start_time) + "\n")
if epoch % save_epoch == (save_epoch - 1):
torch.save(
{
'epoch': epoch + 1,
'state_dict': model.state_dict(),
'opt_dict': optimizer.state_dict(),
},
os.path.join(
save_dir, 'models', modelName + '_' + dataset + '_epoch-' +
str(epoch) + '.pth.tar'))
print("Save model at {}\n".format(
os.path.join(
save_dir, 'models', modelName + '_' + dataset + '_epoch-' +
str(epoch) + '.pth.tar')))
if useTest and epoch % test_interval == (test_interval - 1):
model.eval()
start_time = timeit.default_timer()
running_loss = 0.0
running_corrects = 0
for inputs, labels in tqdm(test_dataloader):
inputs = inputs.to(device)
labels = labels.to(device)
with torch.no_grad():
outputs = model(inputs)
probs = nn.Softmax(dim=1)(outputs)
preds = torch.max(probs, 1)[1]
loss = criterion(outputs, labels)
running_loss += loss.item() * inputs.size(0)
running_corrects += torch.sum(preds == labels.data)
epoch_loss = running_loss / test_size
epoch_acc = running_corrects.double() / test_size
writer.add_scalar('data/test_loss_epoch', epoch_loss, epoch)
writer.add_scalar('data/test_acc_epoch', epoch_acc, epoch)
print("[test] Epoch: {}/{} Loss: {} Acc: {}".format(
epoch + 1, nEpochs, epoch_loss, epoch_acc))
stop_time = timeit.default_timer()
print("Execution time: " + str(stop_time - start_time) + "\n")
writer.close()
def train_model(dataset=dataset, save_dir=SAVE_FILE_FOLDER, num_classes=num_classes, lr=lr,
num_epochs=nEpochs, save_epoch=snapshot, useTest=useTest, test_interval=nTestInterval):
"""
Args:
num_classes (int): Number of classes in the data
num_epochs (int, optional): Number of epochs to train for.
"""
if modelName == 'C3D':
model = C3D_model.C3D(num_classes=num_classes, pretrained=IF_PRETRAIN)
train_params = [{'params': C3D_model.get_1x_lr_params(model), 'lr': lr},
{'params': C3D_model.get_10x_lr_params(model), 'lr': lr * 10}]
elif modelName == 'C3D_td5':
model = C3D_model.C3D_td5(num_classes=num_classes, pretrained=IF_PRETRAIN)
train_params = [{'params': C3D_model.get_1x_lr_params(model), 'lr': lr},
{'params': C3D_model.get_10x_lr_params(model), 'lr': lr * 10}]
elif modelName == 'R2Plus1D':
model = R2Plus1D_model.R2Plus1DClassifier(num_classes=num_classes, layer_sizes=(2, 2, 2, 2))
train_params = [{'params': R2Plus1D_model.get_1x_lr_params(model), 'lr': lr},
{'params': R2Plus1D_model.get_10x_lr_params(model), 'lr': lr * 10}]
elif modelName == 'R3D':
model = R3D_model.R3DClassifier(num_classes=num_classes, layer_sizes=(2, 2, 2, 2))
train_params = model.parameters()
else:
print('We only implemented C3D and R2Plus1D models.')
raise NotImplementedError
criterion = nn.CrossEntropyLoss() # standard crossentropy loss for classification
if _optimizer == "SGD":
optimizer = optim.SGD(train_params, lr=lr, momentum=MOMENTUM, weight_decay=WD)
elif _optimizer == "Adam":
optimizer = optim.Adam(train_params, lr=lr, weight_decay=WD)
# print(optimizer)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=SCHEDULER_STEP_SIZE,
gamma=SCHEDULER_GAMMA) # the scheduler divides the lr by 10 every 10 epochs
model.to(device)
criterion.to(device)
# if resume_epoch == 0:
if resume_model_path == None:
print("Training {} from scratch...".format(modelName))
else:
checkpoint = torch.load(
resume_model_path,
map_location=lambda storage, loc: storage) # Load all tensors onto the CPU
print("Initializing weights from: {}...".format(resume_model_path))
model.load_state_dict(checkpoint['state_dict'])
if RESUM_OPTIMIZER:
optimizer.load_state_dict(checkpoint['opt_dict'])
# resume_epoch
# else:
# checkpoint = torch.load(os.path.join(SAVE_FILE_FOLDER, 'models', EXP_NAME + '_epoch-' + str(resume_epoch - 1) + '.pth.tar'),
# map_location=lambda storage, loc: storage) # Load all tensors onto the CPU
# print("Initializing weights from: {}...".format(
# os.path.join(SAVE_FILE_FOLDER, EXP_NAME + '_epoch-' + str(resume_epoch - 1) + '.pth.tar')))
# model.load_state_dict(checkpoint['state_dict'])
# optimizer.load_state_dict(checkpoint['opt_dict'])
print('Total params: %.2fM' % (sum(p.numel() for p in model.parameters()) / 1000000.0))
model.to(device)
criterion.to(device)
writer = SummaryWriter(logdir=LOG_PATH)
print('Training model on {} dataset...'.format(dataset))
train_dataloader = DataLoader(VideoDataset(dataset=dataset, split='train', clip_len=clip_len, preprocess=IF_PREPROCESS_TRAIN, grayscale=grayscale), batch_size=BS, shuffle=True, num_workers=N_WORKERS)
val_dataloader = DataLoader(VideoDataset(dataset=dataset, split='val', clip_len=clip_len, preprocess=IF_PREPROCESS_VAL, grayscale=grayscale), batch_size=BS, num_workers=N_WORKERS)
test_dataloader = DataLoader(VideoDataset(dataset=dataset, split='test', clip_len=clip_len, preprocess=IF_PREPROCESS_TEST, grayscale=grayscale), batch_size=BS, num_workers=N_WORKERS)
trainval_loaders = {'train': train_dataloader, 'val': val_dataloader}
trainval_sizes = {x: len(trainval_loaders[x].dataset) for x in ['train', 'val']}
test_size = len(test_dataloader.dataset)
cudnn.benchmark = True
global_best_val_acc = 0
for epoch in range(num_epochs):
# each epoch has a training and validation step
for phase in ['train', 'val']:
start_time = timeit.default_timer()
# reset the running loss and corrects
running_loss = 0.0
running_corrects = 0.0
# running_roc = 0.0
list_pred = list()
list_label = list()
# print(optimizer)
# set model to train() or eval() mode depending on whether it is trained
# or being validated. Primarily affects layers such as BatchNorm or Dropout.
if phase == 'train':
# scheduler.step() is to be called once every epoch during training
scheduler.step()
model.train()
else:
model.eval()
# for inputs, labels in tqdm(trainval_loaders[phase]):
run_count = 0
for inputs, labels in trainval_loaders[phase]:
# move inputs and labels to the device the training is taking place on
inputs = Variable(inputs, requires_grad=True).to(device)
labels = Variable(labels).to(device)
optimizer.zero_grad()
if phase == 'train':
outputs = model(inputs)
else:
with torch.no_grad():
outputs = model(inputs)
probs = nn.Softmax(dim=1)(outputs)
preds = torch.max(probs, 1)[1]
loss = criterion(outputs, labels)
if phase == 'train':
loss.backward()
optimizer.step()
running_loss += loss.item() * inputs.size(0)
running_corrects += torch.sum(preds == labels.data)
# try:
# running_roc += roc_auc_score(labels.data.cpu(), preds.cpu())
# except:
# y_true = labels.data.cpu().tolist()
# y_true_2 = y_true.copy()
# for i_cls in range(N_CLASSES):
# y_true_2.append(i_cls)
#
# y_pred = preds.cpu().tolist()
# y_pred_2 = y_pred.copy()
# for i_cls in range(N_CLASSES):
# y_pred_2.append(i_cls)
#
# running_roc += roc_auc_score(y_true_2, y_pred_2)
#
# run_count += 1
list_label += labels.data.cpu().tolist()
list_pred += preds.cpu().tolist()
epoch_loss = running_loss / trainval_sizes[phase]
epoch_acc = running_corrects.double() / trainval_sizes[phase]
epoch_roc = multiclass_roc_score(label=list_label, pred=list_pred, n_cls=N_CLASSES)
if phase == 'train':
writer.add_scalar('data/train_loss_epoch', epoch_loss, epoch)
writer.add_scalar('data/train_acc_epoch', epoch_acc, epoch)
writer.add_scalar('data/train_roc_epoch', epoch_roc, epoch)
else:
writer.add_scalar('data/val_loss_epoch', epoch_loss, epoch)
writer.add_scalar('data/val_acc_epoch', epoch_acc, epoch)
writer.add_scalar('data/val_roc_epoch', epoch_roc, epoch)
# if epoch_acc >= global_best_val_acc:
# torch.save({
# 'epoch': epoch + 1,
# 'state_dict': model.state_dict(),
# 'opt_dict': optimizer.state_dict(),
# }, os.path.join(SAVE_FILE_FOLDER, 'models', EXP_NAME + '_epoch-' + str(epoch) + 'ValAcc_{:10.4f}_'.format(epoch_loss) + '.pth.tar'))
# print("Save model at {}\n".format(
# os.path.join(SAVE_FILE_FOLDER, 'models', EXP_NAME + '_epoch-' + str(epoch) + 'ValAcc_{:10.4f}_'.format(epoch_loss) + '.pth.tar')))
print("[{}] Epoch: {}/{} Loss: {} Acc: {}, ROC:{}".format(phase, epoch+1, nEpochs, epoch_loss, epoch_acc, epoch_roc))
stop_time = timeit.default_timer()
# print("Execution time: " + str(stop_time - start_time) + "\n")
if epoch % save_epoch == (save_epoch - 1):
torch.save({
'epoch': epoch + 1,
'state_dict': model.state_dict(),
'opt_dict': optimizer.state_dict(),
}, os.path.join(SAVE_FILE_FOLDER, EXP_NAME + '_epoch-' + str(epoch) + '.pth.tar'))
print("Save model at {}\n".format(os.path.join(SAVE_FILE_FOLDER, EXP_NAME + '_epoch-' + str(epoch) + '.pth.tar')))
if useTest and epoch % test_interval == (test_interval - 1):
model.eval()
start_time = timeit.default_timer()
running_loss = 0.0
running_corrects = 0.0
# running_roc = 0.0
list_pred = list()
list_label = list()
# for inputs, labels in tqdm(test_dataloader):
run_count = 0
for inputs, labels in test_dataloader:
inputs = inputs.to(device)
labels = labels.to(device)
with torch.no_grad():
outputs = model(inputs)
probs = nn.Softmax(dim=1)(outputs)
preds = torch.max(probs, 1)[1]
loss = criterion(outputs, labels)
running_loss += loss.item() * inputs.size(0)
running_corrects += torch.sum(preds == labels.data)
# try:
# running_roc += roc_auc_score(labels.data.cpu(), preds.cpu())
# except:
# running_roc += 0.5
# run_count += 1
list_label += labels.data.cpu().tolist()
list_pred += preds.cpu().tolist()
epoch_loss = running_loss / test_size
epoch_acc = running_corrects.double() / test_size
epoch_roc = multiclass_roc_score(label=list_label, pred=list_pred, n_cls=N_CLASSES)
writer.add_scalar('data/test_loss_epoch', epoch_loss, epoch)
writer.add_scalar('data/test_acc_epoch', epoch_acc, epoch)
writer.add_scalar('data/test_roc_epoch', epoch_roc, epoch)
print("[test] Epoch: {}/{} Loss: {} Acc:{} ROC: {}".format(epoch+1, nEpochs, epoch_loss, epoch_acc, epoch_roc))
stop_time = timeit.default_timer()
# print("Execution time: " + str(stop_time - start_time) + "\n")
writer.close()
from network import R2Plus1D_model
import torch
import config_net as config
import imageio
import cv2
import numpy as np
import time
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from torch.autograd import Variable
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
net = R2Plus1D_model.R2Plus1DClassifier(2, (2, 2, 2, 2),
pretrained=True).to(device)
video_id = '93.mp4'
filename = config.test_dir + video_id
capture = cv2.VideoCapture(filename)
video_max_len = int(capture.get(cv2.CAP_PROP_FRAME_COUNT))
capture.release()
vid = imageio.get_reader(filename, 'ffmpeg')
test_start = 4 * 60 * config.video_fps
test_end = video_max_len
anomaly_graph = []
f = open('check_video_2.txt', 'w')
start_time = time.time()
for frame_start in range(test_start, test_end, config.test_video_step):
if (frame_start + config.prepare_len > test_end):
def main():
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print("Device being used:", device)
with open('./dataloaders/ferryboat_labels.txt', 'r') as f:
class_names = f.readlines()
f.close()
# init model
num_classes = 4
modelName = 'STP'
if modelName == 'I3D':
model = I3D_model.InceptionI3d(num_classes=num_classes, in_channels=3)
size = (240, 284)
crop_size = 224
elif modelName == 'R2Plus1D':
model = R2Plus1D_model.R2Plus1DClassifier(num_classes=num_classes,
layer_sizes=(3, 4, 6, 3))
size = (171, 128)
crop_size = 112
elif modelName == 'C3D':
model = C3D_model.C3D(num_classes=num_classes, pretrained=False)
size = (171, 128)
crop_size = 112
elif modelName == 'P3D':
model = p3d_model.P3D63(num_classes=num_classes)
size = (176, 210)
crop_size = 160
elif modelName == 'R3D':
model = R3D_model.R3DClassifier(num_classes=num_classes,
layer_sizes=(3, 4, 6, 3))
size = (171, 128)
crop_size = 112
elif modelName == 'STP':
model = STP_model.STP(num_classes=num_classes, in_channels=3)
size = (240, 284)
crop_size = 224
checkpoint = torch.load('./models/I3D-ferryboat4_epoch-199.pth.tar',
map_location=lambda storage, loc: storage)
model_dict = model.state_dict()
checkpoint_load = {
k: v
for k, v in (checkpoint['state_dict']).items() if k in model_dict
}
model_dict.update(checkpoint_load)
model.load_state_dict(model_dict)
model.to(device)
model.eval()
for root, dirs, files in os.walk('./VAR/ferryboat/test/'):
l_names = locals()
l_names['Inshore'] = 0
l_names['Neg'] = 0
l_names['Offshore'] = 0
l_names['Traffic'] = 0
l_names['Inshore1'] = 0
l_names['Neg1'] = 0
l_names['Offshore1'] = 0
l_names['Traffic1'] = 0
if len(dirs) > 4:
video_name = dirs
for name in video_name:
class_name = name.split('_')[1]
video = './ferryboat/' + class_name + "/" + name + '.avi'
clip = []
cap = cv2.VideoCapture(video)
retaining = True
while retaining:
retaining, frame = cap.read()
if not retaining and frame is None:
continue
tmp_ = center_crop(cv2.resize(frame, size), crop_size)
tmp = tmp_
clip.append(tmp)
if len(clip) == 16:
inputs = np.array(clip).astype(np.float32)
inputs = np.expand_dims(inputs, axis=0)
inputs = np.transpose(inputs, (0, 4, 1, 2, 3))
inputs = torch.from_numpy(inputs)
inputs = torch.autograd.Variable(
inputs, requires_grad=False).to(device)
with torch.no_grad():
if modelName == 'STP':
outputs, index = model.forward(inputs)
else:
outputs = model.forward(inputs)
iii = index.cpu().data
probs = torch.nn.Softmax(dim=1)(outputs)
label = torch.max(probs,
1)[1].detach().cpu().numpy()[0]
if modelName == 'I3D':
label = int(label[0])
pre = class_names[label].split(' ')[1][:-1]
l_names[str(class_name)] = l_names[str(class_name)] + 1
if str(pre) == str(class_name):
l_names[str(class_name) +
'1'] = l_names[str(class_name) + '1'] + 1
elif str(class_name) == 'Neg':
l_names[str(class_name) +
'1'] = l_names[str(class_name) + '1'] + 1
elif str(class_name) == 'Traffic':
l_names[str(class_name) +
'1'] = l_names[str(class_name) + '1'] + 1
clip.pop(0)
cv2.waitKey(30)
cap.release()
cv2.destroyAllWindows()
print(
str(class_name) + '_acc:' + str(
int(l_names[str(class_name) + '1']) /
int(l_names[str(class_name)])))