def main_run(dataset, stage, root_dir, out_dir, stage1_dict, seqLen,
trainBatchSize, numEpochs, lr1, decay_factor, decay_step, memSize,
outPool_size, split, evalInterval, regression, rloss, debug):
if debug:
n_workers = 0
n_workers_test = 0
device = 'cpu'
else:
n_workers = 4
n_workers_test = 2
device = 'cuda'
# Train/Validation/Test split
train_splits = ["S1", "S3", "S4"]
val_splits = ["S2"]
test_split = split
mean = [0.485, 0.456, 0.406]
std = [0.229, 0.224, 0.225]
normalize = Normalize(mean=mean, std=std)
stage = stage
#test_split = split
seqLen = seqLen
memSize = memSize
c_cam_classes = outPool_size
dataset = dataset
best_acc = 0
if stage == 1:
trainBatchSize = trainBatchSize
testBatchSize = trainBatchSize
lr1 = lr1
decay_factor = decay_factor
decay_step = decay_step
numEpochs = numEpochs
elif stage == 2 or stage == 3:
trainBatchSize = trainBatchSize
testBatchSize = trainBatchSize
lr1 = lr1
decay_factor = decay_factor
decay_step = decay_step
numEpochs = numEpochs
num_classes = 61
dataset_dir = root_dir
#model_folder = os.path.join('.', out_dir, dataset, str(test_split))
model_folder = os.path.join('./', out_dir, 'stage' + str(stage))
if not os.path.exists(model_folder):
os.makedirs(model_folder)
else:
print('Directory {} exists!'.format(model_folder))
sys.exit()
note_fl = open(model_folder + '/note.txt', 'w')
note_fl.write('Number of Epochs = {}\n'
'lr = {}\n'
'Train Batch Size = {}\n'
'Sequence Length = {}\n'
'Decay steps = {}\n'
'Decay factor = {}\n'
'Memory size = {}\n'
'Memory cam classes = {}\n'.format(numEpochs, lr1,
trainBatchSize, seqLen,
decay_step, decay_factor,
memSize, c_cam_classes))
note_fl.close()
# Log files
writer = SummaryWriter(model_folder)
train_log_loss = open((model_folder + '/train_log_loss.txt'), 'w')
train_log_acc = open((model_folder + '/train_log_acc.txt'), 'w')
train_log_loss_batch = open((model_folder + '/train_log_loss_batch.txt'),
'w')
test_log_loss = open((model_folder + '/test_log_loss.txt'), 'w')
test_log_acc = open((model_folder + '/test_log_acc.txt'), 'w')
# Dataloaders
spatial_transform = Compose([
Scale(256),
RandomHorizontalFlip(),
MultiScaleCornerCrop([1, 0.875, 0.75, 0.65625], 224),
ToTensor(), normalize
])
print('Preparing dataset...')
# Data loader
normalize = Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
spatial_transform = Compose([
Scale(256),
RandomHorizontalFlip(),
MultiScaleCornerCrop([1, 0.875, 0.75, 0.65625], 224),
# ToTensor(),
# normalize
])
transform_rgb = Compose([ToTensor(), normalize])
transform_MS = Compose([Resize((7, 7)), ToTensor()])
vid_seq_train = makeDataset(dataset_dir,
splits=train_splits,
spatial_transform=spatial_transform,
transform_rgb=transform_rgb,
transform_MS=transform_MS,
seqLen=seqLen,
fmt='.png',
regression=regression)
train_loader = torch.utils.data.DataLoader(vid_seq_train,
batch_size=trainBatchSize,
shuffle=True,
num_workers=n_workers,
pin_memory=True)
vid_seq_test = makeDataset(dataset_dir,
splits=val_splits,
spatial_transform=Compose(
[Scale(256), CenterCrop(224)]),
transform_rgb=transform_rgb,
transform_MS=transform_MS,
seqLen=seqLen,
fmt='.png',
regression=regression,
verbose=False)
test_loader = torch.utils.data.DataLoader(vid_seq_test,
batch_size=testBatchSize,
shuffle=False,
num_workers=n_workers,
pin_memory=True)
print('here')
print('Number of train samples = {}'.format(vid_seq_train.__len__()))
print('Number of test samples = {}'.format(vid_seq_test.__len__()))
train_params = []
if stage == 1:
if regression:
model = attentionModel(num_classes=num_classes,
mem_size=memSize,
n_channels=1)
else:
model = attentionModel(num_classes=num_classes, mem_size=memSize)
model.train(False)
for params in model.parameters():
params.requires_grad = False
elif stage == 2 or stage == 3:
if regression:
model = attentionModel(num_classes=num_classes,
mem_size=memSize,
n_channels=1,
c_cam_classes=c_cam_classes)
else:
model = attentionModel(num_classes=num_classes,
mem_size=memSize,
c_cam_classes=c_cam_classes)
#model = attentionModel(num_classes=num_classes, mem_size=memSize, c_cam_classes=c_cam_classes)
if stage == 2:
checkpoint_path = os.path.join(
stage1_dict, 'last_checkpoint_stage' + str(1) + '.pth.tar')
elif stage == 3:
checkpoint_path = os.path.join(
stage1_dict, 'last_checkpoint_stage' + str(2) + '.pth.tar')
if os.path.exists(checkpoint_path):
print('Loading weights from checkpoint file {}'.format(
checkpoint_path))
else:
print('Checkpoint file {} does not exist'.format(checkpoint_path))
sys.exit()
last_checkpoint = torch.load(checkpoint_path)
model.load_state_dict(last_checkpoint['model_state_dict'],
strict=False)
model.train(False)
for params in model.parameters():
params.requires_grad = False
for params in model.resNet.layer4[0].conv1.parameters():
params.requires_grad = True
train_params += [params]
for params in model.resNet.layer4[0].conv2.parameters():
params.requires_grad = True
train_params += [params]
for params in model.resNet.layer4[1].conv1.parameters():
params.requires_grad = True
train_params += [params]
for params in model.resNet.layer4[1].conv2.parameters():
params.requires_grad = True
train_params += [params]
for params in model.resNet.layer4[2].conv1.parameters():
params.requires_grad = True
train_params += [params]
for params in model.resNet.layer4[2].conv2.parameters():
params.requires_grad = True
train_params += [params]
for params in model.resNet.fc.parameters():
params.requires_grad = True
train_params += [params]
# Add params from ms_module
if stage == 2:
for params in model.ms_module.parameters():
params.requires_grad = True
train_params += [params]
for params in model.lsta_cell.parameters():
params.requires_grad = True
train_params += [params]
for params in model.classifier.parameters():
params.requires_grad = True
train_params += [params]
model.classifier.train(True)
model.ms_module.train(True)
model.to(device)
loss_fn = nn.CrossEntropyLoss()
if regression:
if rloss == 'MSE':
# Mean Squared Error loss
loss_ms_fn = nn.MSELoss() # it should work
elif rloss == 'L1':
# L1 loss
loss_ms_fn = nn.L1Loss()
elif rloss == 'SmoothL1':
# Huber Loss or Smooth L1 Loss
loss_ms_fn = nn.SmoothL1Loss()
elif rloss == 'KLdiv':
# Kullback-Leiber Loss
loss_ms_fn = nn.KLDivLoss()
else:
# classification
loss_ms_fn = nn.CrossEntropyLoss() # TODO: check paper Planamente
optimizer_fn = torch.optim.Adam(train_params,
lr=lr1,
weight_decay=5e-4,
eps=1e-4)
optim_scheduler = torch.optim.lr_scheduler.MultiStepLR(
optimizer_fn, milestones=decay_step, gamma=decay_factor)
train_iter = 0
for epoch in range(numEpochs):
#optim_scheduler.step()
epoch_loss = 0
numCorrTrain = 0
trainSamples = 0
iterPerEpoch = 0
# model.classifier.train(True)
model.lsta_cell.train(True)
model.classifier.train(True)
if stage == 2:
model.resNet.layer4[0].conv1.train(True)
model.resNet.layer4[0].conv2.train(True)
model.resNet.layer4[1].conv1.train(True)
model.resNet.layer4[1].conv2.train(True)
model.resNet.layer4[2].conv1.train(True)
model.resNet.layer4[2].conv2.train(True)
model.resNet.fc.train(True)
model.ms_module.train(True)
writer.add_scalar('lr', optimizer_fn.param_groups[0]['lr'], epoch + 1)
for i, (inputs, inputsMS, targets) in enumerate(train_loader):
# Inputs:
# - inputsRGB : the rgb frame input
# Labels :
# - inputsMS : the motion task label
# - targets : output
train_iter += 1
iterPerEpoch += 1
optimizer_fn.zero_grad()
inputVariable = inputs.permute(1, 0, 2, 3, 4).to(device)
labelVariable = targets.to(device)
msVariable = inputsMS.to(device)
trainSamples += inputs.size(0)
output_label, _, output_ms = model(inputVariable, device)
loss_c = loss_fn(output_label, labelVariable)
if stage == 2:
if regression:
msVariable = torch.reshape(
msVariable, (seqLen * 7 * 7, msVariable.size(0)))
output_ms = torch.sigmoid(output_ms)
output_ms = torch.reshape(
output_ms, (seqLen * 7 * 7, output_ms.size(0)))
else:
# classification task
msVariable = torch.reshape(
msVariable,
(seqLen * 7 * 7, msVariable.size(0))).long()
output_ms = torch.reshape(
output_ms, (seqLen * 7 * 7, 2, output_ms.size(0)))
loss_ms = loss_ms_fn(output_ms, msVariable)
loss = loss_c + loss_ms
else:
loss = loss_c
loss.backward()
optimizer_fn.step()
_, predicted = torch.max(output_label.data, 1)
numCorrTrain += (predicted == targets.to(device)).sum()
#print('Training loss after {} iterations = {} '.format(train_iter, loss.data.item()))
#train_log_loss_batch.write('Training loss after {} iterations = {}\n'.format(train_iter, loss.data.item()))
#writer.add_scalar('train/iter_loss', loss.data.item(), train_iter)
epoch_loss += loss.data.item()
avg_loss = epoch_loss / iterPerEpoch
trainAccuracy = (numCorrTrain / trainSamples) * 100
print('Average training loss after {} epoch = {} '.format(
epoch + 1, avg_loss))
print('Training accuracy after {} epoch = {}% '.format(
epoch + 1, trainAccuracy))
writer.add_scalar('train/epoch_loss', avg_loss, epoch + 1)
writer.add_scalar('train/accuracy', trainAccuracy, epoch + 1)
train_log_loss.write('Training loss after {} epoch = {}\n'.format(
epoch + 1, avg_loss))
train_log_acc.write('Training accuracy after {} epoch = {}\n'.format(
epoch + 1, trainAccuracy))
save_path_model = os.path.join(
model_folder, 'last_checkpoint_stage' + str(stage) + '.pth.tar')
save_file = {
'epoch': epoch + 1,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer_fn.state_dict(),
'best_acc': best_acc,
}
torch.save(save_file, save_path_model)
if (epoch + 1) % evalInterval == 0:
#print('Testing...')
model.train(False)
test_loss_epoch = 0
test_iter = 0
test_samples = 0
numCorr = 0
for j, (inputs, inputsMS, targets) in enumerate(test_loader):
#print('testing inst = {}'.format(j))
test_iter += 1
test_samples += inputs.size(0)
inputVariable = inputs.permute(1, 0, 2, 3, 4).to(device)
labelVariable = targets.to(device)
msVariable = inputsMS.to(device)
output_label, _, output_ms = model(inputVariable, device)
test_loss_c = loss_fn(output_label, labelVariable)
if stage == 2:
if regression:
msVariable = torch.reshape(
msVariable, (seqLen * 7 * 7, msVariable.size(0)))
output_ms = torch.sigmoid(output_ms)
output_ms = torch.reshape(
output_ms, (seqLen * 7 * 7, output_ms.size(0)))
else:
# classification task
msVariable = torch.reshape(
msVariable,
(seqLen * 7 * 7, msVariable.size(0))).long()
output_ms = torch.reshape(
output_ms, (seqLen * 7 * 7, 2, output_ms.size(0)))
test_loss_ms = loss_ms_fn(output_ms, msVariable)
test_loss = test_loss_c + test_loss_ms
else:
test_loss = test_loss_c
test_loss_epoch += test_loss.data.item()
_, predicted = torch.max(output_label.data, 1)
numCorr += (predicted == targets.to(device)).sum()
test_accuracy = (numCorr / test_samples) * 100
avg_test_loss = test_loss_epoch / test_iter
print('Test Loss after {} epochs, loss = {}'.format(
epoch + 1, avg_test_loss))
print('Test Accuracy after {} epochs = {}%'.format(
epoch + 1, test_accuracy))
writer.add_scalar('test/epoch_loss', avg_test_loss, epoch + 1)
writer.add_scalar('test/accuracy', test_accuracy, epoch + 1)
test_log_loss.write('Test Loss after {} epochs = {}\n'.format(
epoch + 1, avg_test_loss))
test_log_acc.write('Test Accuracy after {} epochs = {}%\n'.format(
epoch + 1, test_accuracy))
if test_accuracy > best_acc:
best_acc = test_accuracy
save_path_model = os.path.join(
model_folder,
'best_checkpoint_stage' + str(stage) + '.pth.tar')
save_file = {
'epoch': epoch + 1,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer_fn.state_dict(),
'best_acc': best_acc,
}
torch.save(save_file, save_path_model)
optim_scheduler.step()
train_log_loss.close()
train_log_acc.close()
test_log_acc.close()
train_log_loss_batch.close()
test_log_loss.close()
writer.export_scalars_to_json(model_folder + "/all_scalars.json")
writer.close()
def main_run(numEpochs, lr, stepSize, decayRate, trainBatchSize, seqLen,
memSize, evalInterval, evalMode, numWorkers, outDir,
fightsDir_train, noFightsDir_train, fightsDir_test,
noFightsDir_test):
train_dataset_dir_fights = fightsDir_train
train_dataset_dir_noFights = noFightsDir_train
test_dataset_dir_fights = fightsDir_test
test_dataset_dir_noFights = noFightsDir_test
trainDataset, trainLabels, trainNumFrames = make_split(
train_dataset_dir_fights, train_dataset_dir_noFights)
testDataset, testLabels, testNumFrames = make_split(
test_dataset_dir_fights, test_dataset_dir_noFights)
mean = [0.485, 0.456, 0.406]
std = [0.229, 0.224, 0.225]
normalize = Normalize(mean=mean, std=std)
spatial_transform = Compose([
Scale(256),
RandomHorizontalFlip(),
MultiScaleCornerCrop([1, 0.875, 0.75, 0.65625], 224),
ToTensor(), normalize
])
vidSeqTrain = makeDataset(trainDataset,
trainLabels,
trainNumFrames,
spatial_transform=spatial_transform,
seqLen=seqLen)
# trainLoader = torch.utils.data.DataLoader(vidSeqTrain, batch_size=trainBatchSize,
# shuffle=True, num_workers=numWorkers, pin_memory=True, drop_last=True)
trainLoader = torch.utils.data.DataLoader(vidSeqTrain,
batch_size=trainBatchSize,
shuffle=True,
pin_memory=True,
drop_last=True)
if evalMode == 'centerCrop':
test_spatial_transform = Compose(
[Scale(256), CenterCrop(224),
ToTensor(), normalize])
testBatchSize = 1
elif evalMode == 'tenCrops':
test_spatial_transform = Compose(
[Scale(256), TenCrops(size=224, mean=mean, std=std)])
testBatchSize = 1
elif evalMode == 'fiveCrops':
test_spatial_transform = Compose(
[Scale(256), FiveCrops(size=224, mean=mean, std=std)])
testBatchSize = 1
elif evalMode == 'horFlip':
test_spatial_transform = Compose([
Scale(256),
CenterCrop(224),
FlippedImagesTest(mean=mean, std=std)
])
testBatchSize = 1
vidSeqTest = makeDataset(testDataset,
testLabels,
testNumFrames,
seqLen=seqLen,
spatial_transform=test_spatial_transform)
# testLoader = torch.utils.data.DataLoader(vidSeqTest, batch_size=testBatchSize,
# shuffle=False, num_workers=int(numWorkers/2), pin_memory=True)
testLoader = torch.utils.data.DataLoader(vidSeqTest,
batch_size=testBatchSize,
shuffle=False,
pin_memory=True)
numTrainInstances = vidSeqTrain.__len__()
numTestInstances = vidSeqTest.__len__()
print('Number of training samples = {}'.format(numTrainInstances))
print('Number of testing samples = {}'.format(numTestInstances))
modelFolder = './experiments_' + outDir # Dir for saving models and log files
# Create the dir
if os.path.exists(modelFolder):
print(modelFolder + ' exists!!!')
sys.exit()
else:
os.makedirs(modelFolder)
# Log files
writer = SummaryWriter(modelFolder)
trainLogLoss = open((modelFolder + '/trainLogLoss.txt'), 'w')
trainLogAcc = open((modelFolder + '/trainLogAcc.txt'), 'w')
testLogLoss = open((modelFolder + '/testLogLoss.txt'), 'w')
testLogAcc = open((modelFolder + '/testLogAcc.txt'), 'w')
model = ViolenceModel(mem_size=memSize)
trainParams = []
for params in model.parameters():
params.requires_grad = True
trainParams += [params]
model.train(True)
# model.cuda()
lossFn = nn.CrossEntropyLoss()
optimizerFn = torch.optim.RMSprop(trainParams, lr=lr)
optimScheduler = torch.optim.lr_scheduler.StepLR(optimizerFn, stepSize,
decayRate)
minAccuracy = 50
for epoch in range(numEpochs):
if epoch != 0:
optimScheduler.step()
epochLoss = 0
numCorrTrain = 0
iterPerEpoch = 0
model.train(True)
print('Epoch = {}'.format(epoch + 1))
writer.add_scalar('lr', optimizerFn.param_groups[0]['lr'], epoch + 1)
for i, (inputs, targets) in enumerate(trainLoader):
iterPerEpoch += 1
optimizerFn.zero_grad()
# inputVariable1 = Variable(inputs.permute(1, 0, 2, 3, 4).cuda())
inputVariable1 = Variable(inputs.permute(1, 0, 2, 3, 4))
# labelVariable = Variable(targets.cuda())
print(inputVariable1.shape)
labelVariable = Variable(targets)
# print("labelVariable")
# print(labelVariable)
# print("targets")
# print(targets)
outputLabel = model(inputVariable1)
# print("outputs")
# print(outputLabel)
loss = lossFn(outputLabel, labelVariable)
loss.backward()
optimizerFn.step()
outputProb = torch.nn.Softmax(dim=1)(outputLabel)
# print("outputProb.data going into torch.max ")
# print(outputProb.data)
_, predicted = torch.max(outputProb.data, 1)
# numCorrTrain += (predicted == targets.cuda()).sum()
# print("predicted")
# print(predicted)
numCorrTrain += (predicted == targets).sum()
# epochLoss += loss.data[0]
# print(numCorrTrain)
epochLoss += loss.data.item()
avgLoss = epochLoss / iterPerEpoch
trainAccuracy = (numCorrTrain / numTrainInstances) * 100
print('Training: Loss = {} | Accuracy = {}% '.format(
avgLoss, trainAccuracy))
writer.add_scalar('train/epochLoss', avgLoss, epoch + 1)
writer.add_scalar('train/accuracy', trainAccuracy, epoch + 1)
trainLogLoss.write('Training loss after {} epoch = {}\n'.format(
epoch + 1, avgLoss))
trainLogAcc.write('Training accuracy after {} epoch = {}\n'.format(
epoch + 1, trainAccuracy))
if (epoch + 1) % evalInterval == 0:
model.train(False)
print('Evaluating...')
testLossEpoch = 0
testIter = 0
numCorrTest = 0
for j, (inputs, targets) in enumerate(testLoader):
testIter += 1
if evalMode == 'centerCrop':
# inputVariable1 = Variable(inputs.permute(1, 0, 2, 3, 4).cuda(), volatile=True)
inputVariable1 = Variable(inputs.permute(1, 0, 2, 3, 4),
volatile=True)
else:
# inputVariable1 = Variable(inputs[0].cuda(), volatile=True)
inputVariable1 = Variable(inputs[0], volatile=True)
# labelVariable = Variable(targets.cuda(async =True), volatile=True)
labelVariable = Variable(targets)
outputLabel = model(inputVariable1)
outputLabel_mean = torch.mean(outputLabel, 0, True)
testLoss = lossFn(outputLabel_mean, labelVariable)
testLossEpoch += testLoss.data.item()
_, predicted = torch.max(outputLabel_mean.data, 1)
numCorrTest += (predicted == targets[0]).sum()
testAccuracy = (numCorrTest / numTestInstances) * 100
avgTestLoss = testLossEpoch / testIter
print('Testing: Loss = {} | Accuracy = {}% '.format(
avgTestLoss, testAccuracy))
writer.add_scalar('test/epochloss', avgTestLoss, epoch + 1)
writer.add_scalar('test/accuracy', testAccuracy, epoch + 1)
testLogLoss.write('Test Loss after {} epochs = {}\n'.format(
epoch + 1, avgTestLoss))
testLogAcc.write('Test Accuracy after {} epochs = {}%\n'.format(
epoch + 1, testAccuracy))
if testAccuracy > minAccuracy:
savePathClassifier = (modelFolder + '/bestModel.pth')
torch.save(model, savePathClassifier)
minAccuracy = testAccuracy
trainLogAcc.close()
testLogAcc.close()
trainLogLoss.close()
testLogLoss.close()
writer.export_scalars_to_json(modelFolder + "/all_scalars.json")
writer.close()
return True
# transform_test = transforms.Compose([
# transforms.ToTensor(),
# transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
# ])
data_path = args.datapath
seqLen=args.seqLen
testBatchSize=1
trainX, trainY, testX, testY = make_split(data_path)
mean=[0.485, 0.456, 0.406]
std=[0.229, 0.224, 0.225]
normalize = Normalize(mean=mean, std=std)
spatial_transform = Compose([Scale(256), RandomHorizontalFlip(), MultiScaleCornerCrop([1, 0.875, 0.75, 0.65625], 224),
ToTensor(), normalize])
vidSeqTrain = makeDataset(trainX, trainY, spatial_transform=spatial_transform,
seqLen=seqLen)
trainLoader = torch.utils.data.DataLoader(vidSeqTrain, batch_size=args.trainBatchSize,
shuffle=True, num_workers=0)
test_spatial_transform = Compose([Scale(256), CenterCrop(224), FlippedImagesTest(mean=mean, std=std)])
vidSeqTest = makeDataset(testX, testY, seqLen=seqLen,
spatial_transform=test_spatial_transform)
testLoader = torch.utils.data.DataLoader(vidSeqTest, batch_size=testBatchSize,
shuffle=False, num_workers=1)
# trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform_train)
# trainLoader = torch.utils.data.DataLoader(trainset, batch_size=128, shuffle=True, num_workers=2)
def main_run(dataset, root_dir, checkpoint_path, seqLen, testBatchSize, memSize, outPool_size, split):
mean=[0.485, 0.456, 0.406]
std=[0.229, 0.224, 0.225]
normalize = Normalize(mean=mean, std=std)
test_split = split
seqLen = seqLen
memSize = memSize
c_cam_classes = outPool_size
dataset = dataset
testBatchSize = testBatchSize
checkpoint_path = checkpoint_path
if dataset == 'gtea_61':
num_classes = 61
elif dataset == 'gtea_71':
num_classes = 71
elif dataset == 'egtea_gaze+':
num_classes = 106
else:
print('Wrong dataset')
sys.exit()
dataset_dir = os.path.join(root_dir, dataset)
print('Preparing dataset...')
if dataset == 'egtea_gaze+':
trainDatasetF, testDatasetF, trainLabels, testLabels, trainNumFrames, testNumFrames = gen_split_egtea_gazePlus(dataset_dir,
test_split)
else:
trainDatasetF, testDatasetF, trainLabels, testLabels, trainNumFrames, testNumFrames, _ = gen_split(dataset_dir,
test_split)
vid_seq_test = makeDataset(testDatasetF, testLabels, testNumFrames,
spatial_transform=Compose([Scale(256), CenterCrop(224), ToTensor(), normalize]),
fmt='.jpg', seqLen=seqLen)
print('Number of test samples = {}'.format(vid_seq_test.__len__()))
print("Dataset shape: ", len(vid_seq_test.__getitem__(0)), vid_seq_test.__getitem__(0)[0].shape , end='\n\n\n')
test_loader = torch.utils.data.DataLoader(vid_seq_test, batch_size=testBatchSize,
shuffle=False, num_workers=0, pin_memory=True)
model = attentionModel(num_classes=num_classes, mem_size=memSize, c_cam_classes=c_cam_classes)
if os.path.exists(checkpoint_path):
print('Loading weights from checkpoint file {}'.format(checkpoint_path))
else:
print('Checkpoint file {} does not exist'.format(checkpoint_path))
sys.exit()
last_checkpoint = torch.load(checkpoint_path) #, map_location=torch.device('cpu'))
model.load_state_dict(last_checkpoint['model_state_dict'])
model.cuda()
model.train(False)
model.eval()
print('Testing...')
test_iter = 0
test_samples = 0
numCorr = 0
for j, (inputs, targets) in tqdm(enumerate(test_loader)):
test_iter += 1
test_samples += inputs.size(0)
with torch.no_grad():
print(inputs.shape, targets.shape)
inputVariable = Variable(inputs.permute(1, 0, 2, 3, 4).cuda())
output_label, _ = model(inputVariable)
del inputVariable
_, predicted = torch.max(output_label.data, 1)
numCorr += (predicted == targets.cuda()).sum()
test_accuracy = (numCorr.cpu().item() / test_samples) * 100
print('Test Accuracy after = {}%'.format(test_accuracy))
def main_run(version, flowModel, rgbModel, stackSize, seqLen, memSize, trainDatasetDir, outDir,
trainBatchSize, valBatchSize, lr1, numEpochs, decay_step, decay_factor):
num_classes = 61 # gtea61 dataset
model_folder = os.path.join("./", outDir, version)
# Create the dir
print(f"Checking directory {model_folder}")
if os.path.exists(model_folder):
print('Dir {} exists!'.format(model_folder))
sys.exit()
print(f"Creating directory{model_folder}")
os.makedirs(model_folder)
# Log files
print(f"Creating log files")
train_log_loss = open((model_folder + '/train_log_loss.txt'), 'w')
train_log_acc = open((model_folder + '/train_log_acc.txt'), 'w')
val_log_loss = open((model_folder + '/val_log_loss.txt'), 'w')
val_log_acc = open((model_folder + '/val_log_acc.txt'), 'w')
# ImageNet mean and std
mean = [0.485, 0.456, 0.406]
std = [0.229, 0.224, 0.225]
# Train val partitioning
train_usr = ["S1", "S3", "S4"]
val_usr = ["S2"]
normalize = Normalize(mean=mean, std=std)
spatial_transform = Compose([Scale(256), RandomHorizontalFlip(), MultiScaleCornerCrop([1, 0.875, 0.75, 0.65625], 224),
ToTensor(), normalize])
# train dataset
print(f"Defining train dataset")
vid_seq_train = makeDataset(trainDatasetDir, train_usr, spatial_transform,
stackSize=stackSize, seqLen=seqLen)
train_loader = torch.utils.data.DataLoader(vid_seq_train, batch_size=trainBatchSize,
shuffle=True, num_workers=4, pin_memory=True)
# val dataset
print(f"Defining validation dataset")
vid_seq_val = makeDataset(trainDatasetDir, val_usr,
spatial_transform=Compose([Scale(256), CenterCrop(224), ToTensor(), normalize]),
stackSize=stackSize, phase="val", seqLen=seqLen)
val_loader = torch.utils.data.DataLoader(vid_seq_val, batch_size=valBatchSize,
shuffle=False, num_workers=2, pin_memory=True)
valSamples = vid_seq_val.__len__()
# model
print("Building model")
model = twoStreamAttentionModel(flowModel=flowModel, frameModel=rgbModel, stackSize=stackSize, memSize=memSize, # see twoStreamModel.py
num_classes=num_classes)
print("Setting trainable parameters")
for params in model.parameters(): # initially freeze all layers
params.requires_grad = False
model.train(False)
train_params = []
for params in model.classifier.parameters(): # unfreeze classifier layer (the layer that joins the two models outputs)
params.requires_grad = True
train_params += [params]
for params in model.frameModel.lstm_cell.parameters(): # unfreeze lstm layer of the frame model
train_params += [params]
params.requires_grad = True
for params in model.frameModel.resNet.layer4[0].conv1.parameters(): #unfreeze layer 4
params.requires_grad = True
train_params += [params]
for params in model.frameModel.resNet.layer4[0].conv2.parameters():
params.requires_grad = True
train_params += [params]
for params in model.frameModel.resNet.layer4[1].conv1.parameters():
params.requires_grad = True
train_params += [params]
for params in model.frameModel.resNet.layer4[1].conv2.parameters():
params.requires_grad = True
train_params += [params]
for params in model.frameModel.resNet.layer4[2].conv1.parameters():
params.requires_grad = True
train_params += [params]
#
for params in model.frameModel.resNet.layer4[2].conv2.parameters():
params.requires_grad = True
train_params += [params]
#
for params in model.frameModel.resNet.fc.parameters(): # unfreeze last fully connected layer of frame model
params.requires_grad = True # (I still don't know why, because in the joining of the two models, this layer is skipped)
train_params += [params]
base_params = []
for params in model.flowModel.layer4.parameters(): # unfreeze layer 4 of flow model
base_params += [params]
params.requires_grad = True
print("Moving model to GPU")
model.to(DEVICE)
trainSamples = vid_seq_train.__len__()
min_accuracy = 0
print("Defining loss function, optimizer and scheduler")
loss_fn = nn.CrossEntropyLoss() # loss function
optimizer_fn = torch.optim.SGD([ # optimizer
{'params': train_params},
{'params': base_params, 'lr': 1e-4}, # 1e-4
], lr=lr1, momentum=0.9, weight_decay=5e-4)
#scheduler
optim_scheduler = torch.optim.lr_scheduler.StepLR(optimizer_fn, step_size=decay_step, gamma=decay_factor)
train_iter = 0
print("Training begun")
# TRAIN PROCEDURE
for epoch in range(numEpochs):
optim_scheduler.step()
epoch_loss = 0
numCorrTrain = 0
iterPerEpoch = 0
model.classifier.train(True)
model.flowModel.layer4.train(True)
start = time.time()
for j, (inputFrame, inputMMaps, inputFlow, targets) in enumerate(train_loader):
print(f"step {j} / {int(np.floor(trainSamples/trainBatchSize))}")
train_iter += 1
iterPerEpoch += 1
optimizer_fn.zero_grad() # put gradients to zero
inputVariableFlow = Variable(inputFlow.to(DEVICE))
inputVariableFrame = Variable(inputFrame.permute(1, 0, 2, 3, 4).to(DEVICE))
labelVariable = Variable(targets.to(DEVICE))
#print("predict")
output_label = model(inputVariableFlow, inputVariableFrame) # predict
loss = loss_fn(F.log_softmax(output_label, dim=1), labelVariable) # compute loss
#print("backprop")
loss.backward()
optimizer_fn.step()
#print("accuracy")
_, predicted = torch.max(output_label.data, 1)
numCorrTrain += (predicted == targets.to(DEVICE)).sum() # counting number of correct predictions
epoch_loss += loss.data.item()
avg_loss = epoch_loss / iterPerEpoch # computing average per epoch loss
trainAccuracy = (numCorrTrain.item() / trainSamples) * 100
print('Average training loss after {} epoch = {} '.format(epoch + 1, avg_loss))
print('Training accuracy after {} epoch = {}% '.format(epoch + 1, trainAccuracy))
train_log_loss.write('Training loss after {} epoch = {}\n'.format(epoch + 1, avg_loss)) # log file
train_log_acc.write('Training accuracy after {} epoch = {}\n'.format(epoch + 1, trainAccuracy)) # log file
print(f"Elapsed : {time.time()-start}")
# VALIDATION
if (epoch + 1) % 5 == 0:
model.train(False)
val_loss_epoch = 0
val_iter = 0
numCorr = 0
for j, (inputFrame, inputMMaps, inputFlow, targets) in enumerate(val_loader):
if j % 1 == 0:
print(f"step {j} / {int(np.floor(vid_seq_val.__len__()/valBatchSize))}")
val_iter += 1
inputVariableFlow = Variable(inputFlow.to(DEVICE))
inputVariableFrame = Variable(inputFrame.permute(1, 0, 2, 3, 4).to(DEVICE))
labelVariable = Variable(targets.to(DEVICE))
output_label = model(inputVariableFlow, inputVariableFrame)
loss = loss_fn(F.log_softmax(output_label, dim=1), labelVariable)
val_loss_epoch += loss.data.item()
_, predicted = torch.max(output_label.data, 1)
numCorr += (predicted == labelVariable.data).sum()
val_accuracy = (numCorr.item() / valSamples) * 100
avg_val_loss = val_loss_epoch / val_iter
print('Val Loss after {} epochs, loss = {}'.format(epoch + 1, avg_val_loss))
print('Val Accuracy after {} epochs = {}%'.format(epoch + 1, val_accuracy))
val_log_loss.write('Val Loss after {} epochs = {}\n'.format(epoch + 1, avg_val_loss)) # log file
val_log_acc.write('Val Accuracy after {} epochs = {}%\n'.format(epoch + 1, val_accuracy)) # log file
if val_accuracy > min_accuracy:
save_path_model = (model_folder + '/model_twoStream_state_dict.pth') # every epoch, check if the val accuracy is improved, if so, save that model
torch.save(model.state_dict(), save_path_model) # in that way, even if the model overfit, you will get always the best model
min_accuracy = val_accuracy # in this way you don't have to care too much about the number of epochs
train_log_loss.close()
train_log_acc.close()
val_log_acc.close()
val_log_loss.close()
def main_run(dataset, stage, root_dir, out_dir, seqLen, trainBatchSize,
numEpochs, lr1, decay_factor, decay_step, memSize, outPool_size,
split, evalInterval):
mean = [0.485, 0.456, 0.406]
std = [0.229, 0.224, 0.225]
normalize = Normalize(mean=mean, std=std)
stage = stage
test_split = split
seqLen = seqLen
memSize = memSize
c_cam_classes = outPool_size
dataset = dataset
best_acc = 0
if stage == 1:
trainBatchSize = trainBatchSize
testBatchSize = trainBatchSize
lr1 = lr1
decay_factor = decay_factor
decay_step = decay_step
numEpochs = numEpochs
elif stage == 2:
trainBatchSize = trainBatchSize
testBatchSize = trainBatchSize
lr1 = lr1
decay_factor = decay_factor
decay_step = decay_step
numEpochs = numEpochs
if dataset == 'gtea_61':
num_classes = 61
elif dataset == 'gtea_71':
num_classes = 71
elif dataset == 'egtea_gaze+':
num_classes = 106
else:
print('Wrong dataset')
sys.exit()
dataset_dir = os.path.join(root_dir, dataset)
model_folder = os.path.join('.', out_dir, dataset, str(test_split))
if not os.path.exists(model_folder):
os.makedirs(model_folder)
note_fl = open(model_folder + '/note.txt', 'w')
note_fl.write('Number of Epochs = {}\n'
'lr = {}\n'
'Train Batch Size = {}\n'
'Sequence Length = {}\n'
'Decay steps = {}\n'
'Decay factor = {}\n'
'Memory size = {}\n'
'Memory cam classes = {}\n'.format(numEpochs, lr1,
trainBatchSize, seqLen,
decay_step, decay_factor,
memSize, c_cam_classes))
note_fl.close()
# Log files
writer = SummaryWriter(model_folder)
train_log_loss = open((model_folder + '/train_log_loss.txt'), 'w')
train_log_acc = open((model_folder + '/train_log_acc.txt'), 'w')
train_log_loss_batch = open((model_folder + '/train_log_loss_batch.txt'),
'w')
test_log_loss = open((model_folder + '/test_log_loss.txt'), 'w')
test_log_acc = open((model_folder + '/test_log_acc.txt'), 'w')
spatial_transform = Compose([
Scale(256),
RandomHorizontalFlip(),
MultiScaleCornerCrop([1, 0.875, 0.75, 0.65625], 224),
ToTensor(), normalize
])
print('Preparing dataset...')
if dataset == 'egtea_gaze+':
trainDatasetF, testDatasetF, trainLabels, testLabels, trainNumFrames, testNumFrames = gen_split_egtea_gazePlus(
dataset_dir, test_split)
else:
trainDatasetF, testDatasetF, trainLabels, testLabels, trainNumFrames, testNumFrames, _ = gen_split(
dataset_dir, test_split)
vid_seq_train = makeDataset(trainDatasetF,
trainLabels,
trainNumFrames,
spatial_transform=spatial_transform,
fmt='.jpg',
seqLen=seqLen)
print('Number of train samples = {}'.format(vid_seq_train.__len__()))
train_loader = torch.utils.data.DataLoader(vid_seq_train,
batch_size=trainBatchSize,
num_workers=4,
pin_memory=True)
vid_seq_test = makeDataset(testDatasetF,
testLabels,
testNumFrames,
spatial_transform=Compose([
Scale(256),
CenterCrop(224),
ToTensor(), normalize
]),
fmt='.jpg',
seqLen=seqLen)
print('Number of test samples = {}'.format(vid_seq_test.__len__()))
test_loader = torch.utils.data.DataLoader(vid_seq_test,
batch_size=testBatchSize,
shuffle=False,
num_workers=2,
pin_memory=True)
train_params = []
if stage == 1:
model = attentionModel(num_classes=num_classes,
mem_size=memSize,
c_cam_classes=c_cam_classes)
model.train(False)
for params in model.parameters():
params.requires_grad = False
elif stage == 2:
model = attentionModel(num_classes=num_classes,
mem_size=memSize,
c_cam_classes=c_cam_classes)
checkpoint_path = os.path.join(
model_folder, 'last_checkpoint_stage' + str(1) + '.pth.tar')
if os.path.exists(checkpoint_path):
print('Loading weights from checkpoint file {}'.format(
checkpoint_path))
else:
print('Checkpoint file {} does not exist'.format(checkpoint_path))
sys.exit()
last_checkpoint = torch.load(checkpoint_path)
model.load_state_dict(last_checkpoint['model_state_dict'])
model.train(False)
for params in model.parameters():
params.requires_grad = False
for params in model.resNet.layer4[0].conv1.parameters():
params.requires_grad = True
train_params += [params]
for params in model.resNet.layer4[0].conv2.parameters():
params.requires_grad = True
train_params += [params]
for params in model.resNet.layer4[1].conv1.parameters():
params.requires_grad = True
train_params += [params]
for params in model.resNet.layer4[1].conv2.parameters():
params.requires_grad = True
train_params += [params]
for params in model.resNet.layer4[2].conv1.parameters():
params.requires_grad = True
train_params += [params]
for params in model.resNet.layer4[2].conv2.parameters():
params.requires_grad = True
train_params += [params]
for params in model.resNet.fc.parameters():
params.requires_grad = True
train_params += [params]
for params in model.lsta_cell.parameters():
params.requires_grad = True
train_params += [params]
for params in model.classifier.parameters():
params.requires_grad = True
train_params += [params]
model.classifier.train(True)
model.cuda()
loss_fn = nn.CrossEntropyLoss()
optimizer_fn = torch.optim.Adam(train_params,
lr=lr1,
weight_decay=5e-4,
eps=1e-4)
optim_scheduler = torch.optim.lr_scheduler.MultiStepLR(
optimizer_fn, milestones=decay_step, gamma=decay_factor)
train_iter = 0
for epoch in range(numEpochs):
optim_scheduler.step()
epoch_loss = 0
numCorrTrain = 0
trainSamples = 0
iterPerEpoch = 0
model.classifier.train(True)
writer.add_scalar('lr', optimizer_fn.param_groups[0]['lr'], epoch + 1)
for i, (inputs, targets) in enumerate(train_loader):
train_iter += 1
iterPerEpoch += 1
optimizer_fn.zero_grad()
inputVariable = Variable(inputs.permute(1, 0, 2, 3, 4).cuda())
labelVariable = Variable(targets.cuda())
trainSamples += inputs.size(0)
output_label, _ = model(inputVariable)
loss = loss_fn(output_label, labelVariable)
loss.backward()
optimizer_fn.step()
_, predicted = torch.max(output_label.data, 1)
numCorrTrain += (predicted == targets.cuda()).sum()
if train_iter % 10 == 0:
print('Training loss after {} iterations = {} '.format(
train_iter, loss.data[0]))
train_log_loss_batch.write(
'Training loss after {} iterations = {}\n'.format(
train_iter, loss.data[0]))
writer.add_scalar('train/iter_loss', loss.data[0], train_iter)
epoch_loss += loss.data[0]
avg_loss = epoch_loss / iterPerEpoch
trainAccuracy = (numCorrTrain / trainSamples) * 100
print('Average training loss after {} epoch = {} '.format(
epoch + 1, avg_loss))
print('Training accuracy after {} epoch = {}% '.format(
epoch + 1, trainAccuracy))
writer.add_scalar('train/epoch_loss', avg_loss, epoch + 1)
writer.add_scalar('train/accuracy', trainAccuracy, epoch + 1)
train_log_loss.write('Training loss after {} epoch = {}\n'.format(
epoch + 1, avg_loss))
train_log_acc.write('Training accuracy after {} epoch = {}\n'.format(
epoch + 1, trainAccuracy))
save_path_model = os.path.join(
model_folder, 'last_checkpoint_stage' + str(stage) + '.pth.tar')
save_file = {
'epoch': epoch + 1,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer_fn.state_dict(),
'best_acc': best_acc,
}
torch.save(save_file, save_path_model)
if (epoch + 1) % evalInterval == 0:
print('Testing...')
model.train(False)
test_loss_epoch = 0
test_iter = 0
test_samples = 0
numCorr = 0
for j, (inputs, targets) in enumerate(test_loader):
print('testing inst = {}'.format(j))
test_iter += 1
test_samples += inputs.size(0)
inputVariable = Variable(inputs.permute(1, 0, 2, 3, 4).cuda(),
volatile=True)
labelVariable = Variable(targets.cuda(async=True),
volatile=True)
output_label, _ = model(inputVariable)
test_loss = loss_fn(output_label, labelVariable)
test_loss_epoch += test_loss.data[0]
_, predicted = torch.max(output_label.data, 1)
numCorr += (predicted == targets.cuda()).sum()
test_accuracy = (numCorr / test_samples) * 100
avg_test_loss = test_loss_epoch / test_iter
print('Test Loss after {} epochs, loss = {}'.format(
epoch + 1, avg_test_loss))
print('Test Accuracy after {} epochs = {}%'.format(
epoch + 1, test_accuracy))
writer.add_scalar('test/epoch_loss', avg_test_loss, epoch + 1)
writer.add_scalar('test/accuracy', test_accuracy, epoch + 1)
test_log_loss.write('Test Loss after {} epochs = {}\n'.format(
epoch + 1, avg_test_loss))
test_log_acc.write('Test Accuracy after {} epochs = {}%\n'.format(
epoch + 1, test_accuracy))
if test_accuracy > best_acc:
best_acc = test_accuracy
save_path_model = os.path.join(
model_folder,
'best_checkpoint_stage' + str(stage) + '.pth.tar')
save_file = {
'epoch': epoch + 1,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer_fn.state_dict(),
'best_acc': best_acc,
}
torch.save(save_file, save_path_model)
train_log_loss.close()
train_log_acc.close()
test_log_acc.close()
train_log_loss_batch.close()
test_log_loss.close()
writer.export_scalars_to_json(model_folder + "/all_scalars.json")
writer.close()
def main_run(numEpochs, lr, stepSize, decayRate, trainBatchSize, seqLen,
evalInterval, evalMode, numWorkers, outDir, modelUsed, pretrained,
train_test_split, directory, crossValidation, folds):
compDataset, classCount, class_names = make_split(directory)
if crossValidation:
data, label = compDataset
kFoldCrossValid(folds, data, label, numEpochs, evalMode, numWorkers,
lr, stepSize, decayRate, trainBatchSize, seqLen)
else:
(trainDataset,
trainLabels), (validationDataset,
validationLabels), (testDataset,
testLabels) = sampleFromClass(
compDataset, classCount,
train_test_split)
model, accuracy = modelTrain(modelUsed, pretrained, trainDataset,
trainLabels, validationDataset,
validationLabels, numEpochs, evalInterval,
evalMode, outDir, numWorkers, lr,
stepSize, decayRate, trainBatchSize,
seqLen, True)
'''for printing confusion matrix'''
mean = [0.485, 0.456, 0.406]
std = [0.229, 0.224, 0.225]
normalize = Normalize(mean=mean, std=std)
if evalMode == 'centerCrop':
test_spatial_transform = Compose(
[Scale(256),
CenterCrop(224),
ToTensor(), normalize])
elif evalMode == 'tenCrops':
test_spatial_transform = Compose(
[Scale(256),
TenCrops(size=224, mean=mean, std=std)])
elif evalMode == 'fiveCrops':
test_spatial_transform = Compose(
[Scale(256),
FiveCrops(size=224, mean=mean, std=std)])
elif evalMode == 'horFlip':
test_spatial_transform = Compose([
Scale(256),
CenterCrop(224),
FlippedImagesTest(mean=mean, std=std)
])
vidSeqTest = makeDataset(testDataset,
testLabels,
seqLen=seqLen,
spatial_transform=test_spatial_transform)
testLoader = torch.utils.data.DataLoader(vidSeqTest,
batch_size=1,
shuffle=False,
num_workers=int(numWorkers /
2),
pin_memory=True)
numTestInstances = vidSeqTest.__len__()
print('Number of test samples = {}'.format(numTestInstances))
modelFolder = './experiments_' + outDir + '_' + modelUsed + '_' + str(
pretrained) # Dir for saving models and log files
savePathClassifier = (modelFolder + '/bestModel.pth')
torch.save(model.state_dict(), savePathClassifier)
'''running test samples and printing confusion matrix'''
model.train(False)
print('Testing...')
LossEpoch = 0
testIter = 0
pred = None
targ = None
numCorrTest = 0
for j, (inputs, targets) in enumerate(testLoader):
testIter += 1
#if evalMode == 'centerCrop':
if (torch.cuda.is_available()):
inputVariable1 = Variable(inputs.permute(1, 0, 2, 3, 4).cuda(),
requires_grad=False)
labelVariable = Variable(targets.cuda(async=True),
requires_grad=False)
else:
inputVariable1 = Variable(inputs.permute(1, 0, 2, 3, 4),
requires_grad=False)
labelVariable = Variable(targets, requires_grad=False)
# else:
# if(torch.cuda.is_available()):
# inputVariable1 = Variable(inputs[0].permute(1, 0, 2, 3, 4).cuda(), requires_grad=False)
# labelVariable = Variable(targets.cuda(async=True), requires_grad=False)
# else:
# inputVariable1 = Variable(inputs[0].permute(1, 0, 2, 3, 4), requires_grad=False)
# labelVariable = Variable(targets, requires_grad=False)
outputLabel = model(inputVariable1)
outputProb = torch.nn.Softmax(dim=1)(outputLabel)
_, predicted = torch.max(outputProb.data, 1)
if pred is None:
pred = predicted.cpu().numpy()
targ = targets[0].cpu().numpy()
else:
pred = np.append(pred, predicted.cpu().numpy())
targ = np.append(targ, targets[0].cpu().numpy())
# if(torch.cuda.is_available()):
# numCorrTest += (predicted == targets[0].cuda()).sum()
# else:
# numCorrTest += (predicted == targets[0]).sum()
# Compute confusion matrix
cnf_matrix = confusion_matrix(targ, pred)
np.set_printoptions(precision=2)
# Plot non-normalized confusion matrix
plt.figure()
plot_confusion_matrix(cnf_matrix,
classes=class_names,
title='Confusion matrix, without normalization')
plt.savefig(modelFolder + "/no_norm_confusion_matrix.png")
# Plot normalized confusion matrix
plt.figure()
plot_confusion_matrix(cnf_matrix,
classes=class_names,
normalize=True,
title='Normalized confusion matrix')
plt.savefig(modelFolder + "/confusion_matrix.png")
return True
def modelTrain(modelUsed, pretrained, trainDataset, trainLabels,
validationDataset, validationLabels, numEpochs, evalInterval,
evalMode, outDir, numWorkers, lr, stepSize, decayRate,
trainBatchSize, seqLen, plotting):
mean = [0.485, 0.456, 0.406]
std = [0.229, 0.224, 0.225]
normalize = Normalize(mean=mean, std=std)
spatial_transform = Compose([
Scale(256),
RandomHorizontalFlip(),
MultiScaleCornerCrop([1, 0.875, 0.75, 0.65625], 224),
ToTensor(), normalize
])
vidSeqTrain = makeDataset(trainDataset,
trainLabels,
spatial_transform=spatial_transform,
seqLen=seqLen)
# torch iterator to give data in batches of specified size
trainLoader = torch.utils.data.DataLoader(vidSeqTrain,
batch_size=trainBatchSize,
shuffle=True,
num_workers=numWorkers,
pin_memory=True,
drop_last=True)
if evalMode == 'centerCrop':
test_spatial_transform = Compose(
[Scale(256), CenterCrop(224),
ToTensor(), normalize])
elif evalMode == 'tenCrops':
test_spatial_transform = Compose(
[Scale(256), TenCrops(size=224, mean=mean, std=std)])
elif evalMode == 'fiveCrops':
test_spatial_transform = Compose(
[Scale(256), FiveCrops(size=224, mean=mean, std=std)])
elif evalMode == 'horFlip':
test_spatial_transform = Compose([
Scale(256),
CenterCrop(224),
FlippedImagesTest(mean=mean, std=std)
])
vidSeqValid = makeDataset(validationDataset,
validationLabels,
seqLen=seqLen,
spatial_transform=test_spatial_transform)
validationLoader = torch.utils.data.DataLoader(vidSeqValid,
batch_size=1,
shuffle=False,
num_workers=int(numWorkers /
2),
pin_memory=True)
numTrainInstances = vidSeqTrain.__len__()
numValidationInstances = vidSeqValid.__len__()
print('Number of training samples = {}'.format(numTrainInstances))
print('Number of validation samples = {}'.format(numValidationInstances))
modelFolder = './experiments_' + outDir + '_' + modelUsed + '_' + str(
pretrained) # Dir for saving models and log files
# Create the dir
if os.path.exists(modelFolder):
pass
else:
os.makedirs(modelFolder)
# Log files
writer = SummaryWriter(modelFolder)
trainLogLoss = open((modelFolder + '/trainLogLoss.txt'), 'a')
trainLogAcc = open((modelFolder + '/trainLogAcc.txt'), 'a')
validationLogLoss = open((modelFolder + '/validLogLoss.txt'), 'a')
validationLogAcc = open((modelFolder + '/validLogAcc.txt'), 'a')
model = ViolenceModel(modelUsed, pretrained)
trainParams = []
for params in model.parameters():
if params.requires_grad:
trainParams += [params]
model.train(True)
if (torch.cuda.is_available()):
model.cuda()
lossFn = nn.CrossEntropyLoss()
optimizerFn = torch.optim.RMSprop(trainParams, lr=lr)
optimizerFn.zero_grad()
optimScheduler = torch.optim.lr_scheduler.StepLR(optimizerFn, stepSize,
decayRate)
minAccuracy = 50
train_loss = []
val_loss = []
train_acc = []
val_acc = []
bestmodel = None
for epoch in range(numEpochs):
optimScheduler.step()
epochLoss = 0
numCorrTrain = 0
iterPerEpoch = 0
model.train(True)
print('Epoch = {}'.format(epoch + 1))
writer.add_scalar('lr', optimizerFn.param_groups[0]['lr'], epoch + 1)
for i, (inputs, targets) in enumerate(trainLoader):
iterPerEpoch += 1
optimizerFn.zero_grad()
if (torch.cuda.is_available()):
inputVariable1 = Variable(inputs.permute(1, 0, 2, 3, 4).cuda())
labelVariable = Variable(targets.cuda())
else:
inputVariable1 = Variable(inputs.permute(1, 0, 2, 3, 4))
labelVariable = Variable(targets)
outputLabel = model(inputVariable1)
loss = lossFn(outputLabel, labelVariable)
loss.backward()
optimizerFn.step()
outputProb = torch.nn.Softmax(dim=1)(outputLabel)
_, predicted = torch.max(outputProb.data, 1)
if (torch.cuda.is_available()):
numCorrTrain += (predicted == targets.cuda()).sum()
else:
numCorrTrain += (predicted == targets).sum()
epochLoss += loss.item()
avgLoss = epochLoss / iterPerEpoch
trainAccuracy = (float(numCorrTrain) * 100) / float(numTrainInstances)
train_loss.append(avgLoss)
train_acc.append(trainAccuracy)
print('Training: Loss = {} | Accuracy = {}% '.format(
avgLoss, trainAccuracy))
writer.add_scalar('train/epochLoss', avgLoss, epoch + 1)
writer.add_scalar('train/accuracy', trainAccuracy, epoch + 1)
trainLogLoss.write('Training loss after {} epoch = {}\n'.format(
epoch + 1, avgLoss))
trainLogAcc.write('Training accuracy after {} epoch = {}\n'.format(
epoch + 1, trainAccuracy))
if (epoch + 1) % evalInterval == 0:
model.train(False)
print('Evaluating...')
validationLossEpoch = 0
validationIter = 0
numCorrTest = 0
for j, (inputs, targets) in enumerate(validationLoader):
validationIter += 1
#if evalMode == 'centerCrop':
if (torch.cuda.is_available()):
inputVariable1 = Variable(inputs.permute(1, 0, 2, 3,
4).cuda(),
requires_grad=False)
labelVariable = Variable(targets.cuda(async=True),
requires_grad=False)
else:
inputVariable1 = Variable(inputs.permute(1, 0, 2, 3, 4),
requires_grad=False)
labelVariable = Variable(targets, requires_grad=False)
# else:
# if(torch.cuda.is_available()):
# inputVariable1 = Variable(inputs[0].permute(1, 0, 2, 3, 4).cuda(), requires_grad=False)
# labelVariable = Variable(targets.cuda(async=True), requires_grad=False)
# else:
# inputVariable1 = Variable(inputs[0].permute(1, 0, 2, 3, 4), requires_grad=False)
# labelVariable = Variable(targets, requires_grad=False)
outputLabel = model(inputVariable1)
validationLoss = lossFn(outputLabel, labelVariable)
validationLossEpoch += validationLoss.item()
outputProb = torch.nn.Softmax(dim=1)(outputLabel)
_, predicted = torch.max(outputProb.data, 1)
if (torch.cuda.is_available()):
numCorrTest += (predicted == targets[0].cuda()).sum()
else:
numCorrTest += (predicted == targets[0]).sum()
validationAccuracy = (float(numCorrTest) *
100) / float(numValidationInstances)
avgValidationLoss = validationLossEpoch / validationIter
val_loss.append(avgValidationLoss)
val_acc.append(validationAccuracy)
print('Testing: Loss = {} | Accuracy = {}% '.format(
avgValidationLoss, validationAccuracy))
writer.add_scalar('test/epochloss', avgValidationLoss, epoch + 1)
writer.add_scalar('test/accuracy', validationAccuracy, epoch + 1)
validationLogLoss.write('valid Loss after {} epochs = {}\n'.format(
epoch + 1, avgValidationLoss))
validationLogAcc.write(
'valid Accuracy after {} epochs = {}%\n'.format(
epoch + 1, validationAccuracy))
if validationAccuracy > minAccuracy:
bestmodel = model
minAccuracy = validationAccuracy
'''plotting the accuracy and loss curves'''
if plotting:
xc = range(1, numEpochs + 1)
xv = []
for i in xc:
if (i % evalInterval == 0):
xv.append(i)
plt.figure(1, figsize=(7, 5))
plt.plot(xc, train_loss)
plt.plot(xv, val_loss)
plt.xlabel('num of Epochs')
plt.ylabel('loss')
plt.title('train_loss vs val_loss')
plt.grid(True)
plt.legend(['train', 'val'])
#print plt.style.available # use bmh, classic,ggplot for big pictures
plt.style.use(['classic'])
plt.savefig(modelFolder + "/lossCurve.png")
plt.figure(2, figsize=(7, 5))
plt.plot(xc, train_acc)
plt.plot(xv, val_acc)
plt.xlabel('num of Epochs')
plt.ylabel('accuracy')
plt.title('train_acc vs val_acc')
plt.grid(True)
plt.legend(['train', 'val'], loc=4)
#print plt.style.available # use bmh, classic,ggplot for big pictures
plt.style.use(['classic'])
plt.savefig(modelFolder + "/accuracyCurve.png")
#plt.show()
trainLogAcc.close()
validationLogAcc.close()
trainLogLoss.close()
validationLogLoss.close()
writer.export_scalars_to_json(modelFolder + "/all_scalars.json")
writer.close()
return bestmodel, validationAccuracy
DEVICE = 'cuda'
#dict_path = "/content/gdrive/My Drive/FINAL_LOGS/two_stream_300flow_16frames/model_twoStream_state_dict.pth"
dict_path = "/content/gdrive/My Drive/Lorenzo/ego-rnn-two-in-one/results/rgb_16frames_two_in_one_no_ms/model_rgb_state_dict.pth"
#dict_path = "/content/gdrive/My Drive/FINAL_LOGS/200+150epochs_RGB_16frames/test_2/model_rgb_state_dict.pth"
#dict_path = "/content/gdrive/My Drive/Lorenzo/ego-rnn-ss-task/rgb_16frames_regression_kl/model_rgb_state_dict.pth"
#dict_path = "/content/gdrive/My Drive/FINAL_LOGS/flow300/model_flow_state_dict.pth"
normalize = Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
spatial_transform = Compose(
[Scale(256), CenterCrop(224),
ToTensor(), normalize])
dataset = makeDataset("/content", ["S2"],
spatial_transform=spatial_transform,
seqLen=16,
phase="train")
loader = torch.utils.data.DataLoader(dataset,
batch_size=1,
shuffle=False,
num_workers=2,
pin_memory=True)
#model = twoStreamAttentionModel(flowModel="", frameModel="", stackSize=5, memSize=512, num_classes=61)
model = attentionModel(num_classes=61, mem_size=512)
#model = flow_resnet34(True, channels=2 * 5, num_classes=61)
model.load_state_dict(torch.load(dict_path), strict=True)
model.to(DEVICE)
model.train(False)
def main_run(dataset, stage, root_dir, out_dir, seqLen, trainBatchSize,
numEpochs, lr1, decay_factor, decay_step, memSize, outPool_size,
split, evalInterval, debug):
if debug:
n_workers = 0
n_workers_test = 0
device = 'cpu'
else:
n_workers = 4
n_workers_test = 2
device = 'cuda'
# Train/Validation/Test split
train_splits = ["S1", "S3", "S4"]
val_splits = ["S2"]
test_split = split
mean = [0.485, 0.456, 0.406]
std = [0.229, 0.224, 0.225]
normalize = Normalize(mean=mean, std=std)
stage = stage
#test_split = split
seqLen = seqLen
memSize = memSize
c_cam_classes = outPool_size
dataset = dataset
best_acc = 0
if stage == 1:
trainBatchSize = trainBatchSize
testBatchSize = trainBatchSize
lr1 = lr1
decay_factor = decay_factor
decay_step = decay_step
numEpochs = numEpochs
elif stage == 2:
trainBatchSize = trainBatchSize
testBatchSize = trainBatchSize
lr1 = lr1
decay_factor = decay_factor
decay_step = decay_step
numEpochs = numEpochs
num_classes = 61
dataset_dir = root_dir
#model_folder = os.path.join('.', out_dir, dataset, str(test_split))
model_folder = os.path.join('./', out_dir, dataset, 'LSTA',
'stage' + str(stage))
if not os.path.exists(model_folder):
os.makedirs(model_folder)
else:
print('Directory {} exists!'.format(model_folder))
sys.exit()
note_fl = open(model_folder + '/note.txt', 'w')
note_fl.write('Number of Epochs = {}\n'
'lr = {}\n'
'Train Batch Size = {}\n'
'Sequence Length = {}\n'
'Decay steps = {}\n'
'Decay factor = {}\n'
'Memory size = {}\n'
'Memory cam classes = {}\n'.format(numEpochs, lr1,
trainBatchSize, seqLen,
decay_step, decay_factor,
memSize, c_cam_classes))
note_fl.close()
# Log files
writer = SummaryWriter(model_folder)
train_log_loss = open((model_folder + '/train_log_loss.txt'), 'w')
train_log_acc = open((model_folder + '/train_log_acc.txt'), 'w')
train_log_loss_batch = open((model_folder + '/train_log_loss_batch.txt'),
'w')
test_log_loss = open((model_folder + '/test_log_loss.txt'), 'w')
test_log_acc = open((model_folder + '/test_log_acc.txt'), 'w')
spatial_transform = Compose([
Scale(256),
RandomHorizontalFlip(),
MultiScaleCornerCrop([1, 0.875, 0.75, 0.65625], 224),
ToTensor(), normalize
])
vid_seq_train = makeDataset(root_dir,
train_splits,
spatial_transform=spatial_transform,
sequence=False,
numSeg=1,
fmt='.png',
seqLen=seqLen)
train_loader = torch.utils.data.DataLoader(vid_seq_train,
batch_size=trainBatchSize,
shuffle=True,
num_workers=n_workers,
pin_memory=True)
vid_seq_test = makeDataset(root_dir,
val_splits,
spatial_transform=Compose([
Scale(256),
CenterCrop(224),
ToTensor(), normalize
]),
sequence=False,
numSeg=1,
fmt='.png',
phase='test',
seqLen=seqLen)
test_loader = torch.utils.data.DataLoader(vid_seq_test,
batch_size=testBatchSize,
shuffle=False,
num_workers=n_workers_test,
pin_memory=True)
print('Number of train samples = {}'.format(vid_seq_train.__len__()))
print('Number of test samples = {}'.format(vid_seq_test.__len__()))
train_params = []
if stage == 1:
model = attentionModel(num_classes=num_classes,
mem_size=memSize,
c_cam_classes=c_cam_classes)
model.train(False)
for params in model.parameters():
params.requires_grad = False
elif stage == 2:
model = attentionModel(num_classes=num_classes,
mem_size=memSize,
c_cam_classes=c_cam_classes)
checkpoint_path = os.path.join(
model_folder, 'last_checkpoint_stage' + str(1) + '.pth.tar')
if os.path.exists(checkpoint_path):
print('Loading weights from checkpoint file {}'.format(
checkpoint_path))
else:
print('Checkpoint file {} does not exist'.format(checkpoint_path))
sys.exit()
last_checkpoint = torch.load(checkpoint_path)
model.load_state_dict(last_checkpoint['model_state_dict'])
model.train(False)
for params in model.parameters():
params.requires_grad = False
for params in model.resNet.layer4[0].conv1.parameters():
params.requires_grad = True
train_params += [params]
for params in model.resNet.layer4[0].conv2.parameters():
params.requires_grad = True
train_params += [params]
for params in model.resNet.layer4[1].conv1.parameters():
params.requires_grad = True
train_params += [params]
for params in model.resNet.layer4[1].conv2.parameters():
params.requires_grad = True
train_params += [params]
for params in model.resNet.layer4[2].conv1.parameters():
params.requires_grad = True
train_params += [params]
for params in model.resNet.layer4[2].conv2.parameters():
params.requires_grad = True
train_params += [params]
for params in model.resNet.fc.parameters():
params.requires_grad = True
train_params += [params]
# 2in1 stream
for params in model.resNet.cond.parameters():
params.requires_grad = True
train_params += [params]
# motion modulation layer
for params in model.resNet.sft.parameters():
params.requires_grad = True
train_params += [params]
for params in model.lsta_cell.parameters():
params.requires_grad = True
train_params += [params]
for params in model.classifier.parameters():
params.requires_grad = True
train_params += [params]
model.classifier.train(True)
model.to(device)
loss_fn = nn.CrossEntropyLoss()
optimizer_fn = torch.optim.Adam(train_params,
lr=lr1,
weight_decay=5e-4,
eps=1e-4)
optim_scheduler = torch.optim.lr_scheduler.MultiStepLR(
optimizer_fn, milestones=decay_step, gamma=decay_factor)
train_iter = 0
for epoch in range(numEpochs):
#optim_scheduler.step()
epoch_loss = 0
numCorrTrain = 0
trainSamples = 0
iterPerEpoch = 0
model.classifier.train(True)
writer.add_scalar('lr', optimizer_fn.param_groups[0]['lr'], epoch + 1)
for i, (inputFlow, inputRGB, targets) in enumerate(train_loader):
train_iter += 1
iterPerEpoch += 1
optimizer_fn.zero_grad()
inputVariable = inputRGB.permute(1, 0, 2, 3, 4).to(device)
inputFlow = inputFlow.view(
(inputFlow.shape[0], int(inputFlow.shape[1] / 2), 2,
inputFlow.shape[2], inputFlow.shape[3]))
inputFlow = inputFlow.permute(1, 0, 2, 3, 4).to(device)
labelVariable = targets.to(device)
trainSamples += inputRGB.size(0)
output_label, _ = model(inputVariable, inputFlow, device)
loss = loss_fn(output_label, labelVariable)
loss.backward()
optimizer_fn.step()
_, predicted = torch.max(output_label.data, 1)
numCorrTrain += (predicted == targets.to(device)).sum()
#print('Training loss after {} iterations = {} '.format(train_iter, loss.data.item()))
#train_log_loss_batch.write('Training loss after {} iterations = {}\n'.format(train_iter, loss.data.item()))
#writer.add_scalar('train/iter_loss', loss.data.item(), train_iter)
epoch_loss += loss.data.item()
avg_loss = epoch_loss / iterPerEpoch
trainAccuracy = (numCorrTrain / trainSamples) * 100
print('Average training loss after {} epoch = {} '.format(
epoch + 1, avg_loss))
print('Training accuracy after {} epoch = {}% '.format(
epoch + 1, trainAccuracy))
writer.add_scalar('train/epoch_loss', avg_loss, epoch + 1)
writer.add_scalar('train/accuracy', trainAccuracy, epoch + 1)
train_log_loss.write('Training loss after {} epoch = {}\n'.format(
epoch + 1, avg_loss))
train_log_acc.write('Training accuracy after {} epoch = {}\n'.format(
epoch + 1, trainAccuracy))
save_path_model = os.path.join(
model_folder, 'last_checkpoint_stage' + str(stage) + '.pth.tar')
save_file = {
'epoch': epoch + 1,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer_fn.state_dict(),
'best_acc': best_acc,
}
torch.save(save_file, save_path_model)
if (epoch + 1) % evalInterval == 0:
#print('Testing...')
model.train(False)
test_loss_epoch = 0
test_iter = 0
test_samples = 0
numCorr = 0
for j, (inputFlow, inputRGB, targets) in enumerate(test_loader):
#print('testing inst = {}'.format(j))
test_iter += 1
test_samples += inputRGB.size(0)
inputVariable = inputRGB.permute(1, 0, 2, 3, 4).to(device)
inputFlow = inputFlow.view(
(inputFlow.shape[0], int(inputFlow.shape[1] / 2), 2,
inputFlow.shape[2], inputFlow.shape[3]))
inputFlow = inputFlow.permute(1, 0, 2, 3, 4).to(device)
labelVariable = targets.to(device)
output_label, _ = model(inputVariable, inputFlow, device)
test_loss = loss_fn(output_label, labelVariable)
test_loss_epoch += test_loss.data.item()
_, predicted = torch.max(output_label.data, 1)
numCorr += (predicted == targets.to(device)).sum()
test_accuracy = (numCorr / test_samples) * 100
avg_test_loss = test_loss_epoch / test_iter
print('Test Loss after {} epochs, loss = {}'.format(
epoch + 1, avg_test_loss))
print('Test Accuracy after {} epochs = {}%'.format(
epoch + 1, test_accuracy))
writer.add_scalar('test/epoch_loss', avg_test_loss, epoch + 1)
writer.add_scalar('test/accuracy', test_accuracy, epoch + 1)
test_log_loss.write('Test Loss after {} epochs = {}\n'.format(
epoch + 1, avg_test_loss))
test_log_acc.write('Test Accuracy after {} epochs = {}%\n'.format(
epoch + 1, test_accuracy))
if test_accuracy > best_acc:
best_acc = test_accuracy
save_path_model = os.path.join(
model_folder,
'best_checkpoint_stage' + str(stage) + '.pth.tar')
save_file = {
'epoch': epoch + 1,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer_fn.state_dict(),
'best_acc': best_acc,
}
torch.save(save_file, save_path_model)
optim_scheduler.step()
train_log_loss.close()
train_log_acc.close()
test_log_acc.close()
train_log_loss_batch.close()
test_log_loss.close()
writer.export_scalars_to_json(model_folder + "/all_scalars.json")
writer.close()
def main_run(version, stage, train_data_dir, stage1_dict, out_dir, seqLen,
trainBatchSize, valBatchSize, numEpochs, lr1, decay_factor,
decay_step, mem_size):
num_classes = 61
model_folder = os.path.join("./", out_dir, version)
if os.path.exists(model_folder):
print('Directory {} exists!'.format(model_folder))
sys.exit()
os.makedirs(model_folder)
train_log_loss = open((model_folder + '/train_log_loss.txt'), 'w')
train_log_acc = open((model_folder + '/train_log_acc.txt'), 'w')
val_log_loss = open((model_folder + '/val_log_loss.txt'), 'w')
val_log_acc = open((model_folder + '/val_log_acc.txt'), 'w')
# Train val partitioning
train_usr = ["S1", "S3", "S4"]
val_usr = ["S2"]
# Data loader
normalize = Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
spatial_transform = Compose([
Scale(256),
RandomHorizontalFlip(),
MultiScaleCornerCrop([1, 0.875, 0.75, 0.65625], 224),
ToTensor(), normalize
])
vid_seq_train = makeDataset(train_data_dir,
train_usr,
spatial_transform=spatial_transform,
seqLen=seqLen)
train_loader = torch.utils.data.DataLoader(vid_seq_train,
batch_size=trainBatchSize,
shuffle=True,
num_workers=4,
pin_memory=True)
vid_seq_val = makeDataset(train_data_dir,
val_usr,
spatial_transform=Compose([
Scale(256),
CenterCrop(224),
ToTensor(), normalize
]),
seqLen=seqLen,
phase="test")
val_loader = torch.utils.data.DataLoader(vid_seq_val,
batch_size=valBatchSize,
shuffle=False,
num_workers=2,
pin_memory=True)
train_params = []
# stage 1: train only lstm
if stage == 1:
model = attentionModel(num_classes=num_classes, mem_size=mem_size)
model.train(False)
for params in model.parameters():
params.requires_grad = False
# stage 2: train lstm, layer4, spatial attention and final fc
if stage > 1:
model = attentionModel(num_classes=num_classes, mem_size=mem_size)
model.load_state_dict(torch.load(stage1_dict),
strict=False) # pretrained
model.train(False)
for params in model.parameters():
params.requires_grad = False
#
for params in model.resNet.layer4[0].conv1.parameters():
params.requires_grad = True
train_params += [params]
for params in model.resNet.layer4[0].conv2.parameters():
params.requires_grad = True
train_params += [params]
for params in model.resNet.layer4[1].conv1.parameters():
params.requires_grad = True
train_params += [params]
for params in model.resNet.layer4[1].conv2.parameters():
params.requires_grad = True
train_params += [params]
for params in model.resNet.layer4[2].conv1.parameters():
params.requires_grad = True
train_params += [params]
#
for params in model.resNet.layer4[2].conv2.parameters():
params.requires_grad = True
train_params += [params]
#
for params in model.resNet.fc.parameters(): # fully connected layer
params.requires_grad = True
train_params += [params]
model.resNet.layer4[0].conv1.train(True)
model.resNet.layer4[0].conv2.train(True)
model.resNet.layer4[1].conv1.train(True)
model.resNet.layer4[1].conv2.train(True)
model.resNet.layer4[2].conv1.train(True)
model.resNet.layer4[2].conv2.train(True)
model.resNet.fc.train(True)
if stage == 2:
# set to train the self supervised
for params in model.ms_task.parameters():
params.requires_grad = True
train_params += [params]
if stage == 3:
# motion condition layer
for params in model.resNet.condNet.parameters():
params.requires_grad = True
train_params += [params]
# motion modulation layer
for params in model.resNet.sft.parameters():
params.requires_grad = True
train_params += [params]
# first convolution unfreeze
for params in model.resNet.conv1.parameters():
params.requires_grad = True
train_params += [params]
for params in model.lstm_cell.parameters(
): # for both stages we train the lstm
params.requires_grad = True
train_params += [params]
for params in model.classifier.parameters(
): # for both stages we train the last classifier (after the lstm and avg pooling)
params.requires_grad = True
train_params += [params]
model.lstm_cell.train(True)
model.classifier.train(True)
model.cuda()
loss_fn = nn.CrossEntropyLoss()
optimizer_fn = torch.optim.Adam(train_params,
lr=lr1,
weight_decay=4e-5,
eps=1e-4)
optim_scheduler = torch.optim.lr_scheduler.MultiStepLR(
optimizer_fn, milestones=decay_step, gamma=decay_factor)
train_iter = 0
min_accuracy = 0
for epoch in range(numEpochs):
optim_scheduler.step()
epoch_loss = 0
epoch_ms_loss = 0
numCorrTrain = 0
trainSamples = 0
iterPerEpoch = 0
model.lstm_cell.train(True)
model.classifier.train(True)
if stage > 1:
model.resNet.layer4[0].conv1.train(True)
model.resNet.layer4[0].conv2.train(True)
model.resNet.layer4[1].conv1.train(True)
model.resNet.layer4[1].conv2.train(True)
model.resNet.layer4[2].conv1.train(True)
model.resNet.layer4[2].conv2.train(True)
model.resNet.fc.train(True)
if stage == 2:
model.ms_task.conv.train(True)
model.ms_task.fc.train(True)
if stage == 3:
# motion layers
model.resNet.sft.train(True)
model.resNet.condNet.train(True)
model.resNet.conv1.train(True)
start = time.time()
for i, (inputFrame, inputMMaps, inputFlow,
targets) in enumerate(train_loader):
train_iter += 1
iterPerEpoch += 1
optimizer_fn.zero_grad()
inputFrameVariable = Variable(
inputFrame.permute(
1, 0, 2, 3,
4).to(DEVICE)) # sequence length as first dimension
# flow data is reshaped to have x and y as channels
inputFlow = inputFlow.view(
(inputFlow.shape[0], int(inputFlow.shape[1] / 2), 2,
inputFlow.shape[2], inputFlow.shape[3]))
inputFlowVariable = Variable(
inputFlow.permute(
1, 0, 2, 3,
4).to(DEVICE)) # sequence length as first dimension
labelVariable = Variable(targets.to(DEVICE))
trainSamples += inputFrame.size(0)
output_label, _, ms_lab = model(
(inputFrameVariable, inputFlowVariable), stage)
loss = loss_fn(output_label, labelVariable)
if stage == 2:
ms_loss = loss_ms_fn(F.softmax(ms_lab, dim=1),
mmaps_preparation(inputMMaps, 1e-2))
epoch_ms_loss += ms_loss.data.item() #
(loss + ms_loss).backward()
else:
loss.backward()
optimizer_fn.step()
_, predicted = torch.max(output_label.data, 1)
numCorrTrain += (predicted == targets.to(
DEVICE)).sum() # evaluating number of correct classifications
epoch_loss += loss.data.item()
print(f"Elapsed {time.time() - start}")
avg_loss = epoch_loss / iterPerEpoch
if stage == 2:
avg_ms_loss = epoch_ms_loss / iterPerEpoch #
trainAccuracy = (numCorrTrain.data.item() / trainSamples)
train_log_loss.write('Training loss after {} epoch = {}\n'.format(
epoch + 1, avg_loss)) # log file
train_log_acc.write('Training accuracy after {} epoch = {}\n'.format(
epoch + 1, trainAccuracy)) # log file
print('Train: Epoch = {} | Loss = {} | Accuracy = {}'.format(
epoch + 1, avg_loss, trainAccuracy))
if stage == 2:
train_log_loss.write(
'Training ms loss after {} epoch = {}\n'.format(
epoch + 1, avg_ms_loss)) # log file
print("ms avg epoch loss: {}".format(avg_ms_loss))
if (epoch + 1) % VAL_FREQUENCY == 0:
model.train(False)
val_loss_epoch = 0
val_ms_loss_epoch = 0
val_iter = 0
val_samples = 0
numCorr = 0
for j, (inputFrame, inputMMaps, inputFlow,
targets) in enumerate(val_loader):
val_iter += 1
val_samples += inputFrame.size(0)
inputFrameVariable = Variable(
inputFrame.permute(1, 0, 2, 3, 4).to(DEVICE))
# flow data is reshaped to have x and y as channels
inputFlow = inputFlow.view(
(inputFlow.shape[0], int(inputFlow.shape[1] / 2), 2,
inputFlow.shape[2], inputFlow.shape[3]))
inputFlowVariable = Variable(
inputFlow.permute(
1, 0, 2, 3,
4).to(DEVICE)) # sequence length as first dimension
labelVariable = Variable(targets.to(DEVICE))
output_label, _, ms_lab = model(
(inputFrameVariable, inputFlowVariable), stage)
val_loss = loss_fn(output_label, labelVariable)
val_loss_epoch += val_loss.data.item()
if stage == 2:
ms_loss = loss_ms_fn(F.softmax(ms_lab, dim=1),
mmaps_preparation(inputMMaps,
1e-2)) #
val_ms_loss_epoch += ms_loss.data.item() #
_, predicted = torch.max(output_label.data, 1)
numCorr += (predicted == targets.to(DEVICE)).sum(
) # evaluating number of correct classifications
val_accuracy = (numCorr.data.item() / val_samples)
avg_val_loss = val_loss_epoch / val_iter
avg_ms_loss = val_ms_loss_epoch / val_iter #
print('Val: Epoch = {} | Loss {} | Accuracy = {}'.format(
epoch + 1, avg_val_loss, val_accuracy))
val_log_loss.write('Val Loss after {} epochs = {}\n'.format(
epoch + 1, avg_val_loss)) # log file
val_log_acc.write('Val Accuracy after {} epochs = {}%\n'.format(
epoch + 1, val_accuracy)) # log file
if stage == 2:
val_log_loss.write('Val ms loss after {} epoch = {}\n'.format(
epoch + 1, avg_ms_loss)) # log file
print("Ms Epoch: {}".format(avg_ms_loss))
if val_accuracy > min_accuracy:
save_path_model = (
model_folder + '/model_rgb_state_dict.pth'
) # every epoch, check if the val accuracy is improved, if so, save that model
torch.save(
model.state_dict(), save_path_model
) # in that way, even if the model overfit, you will get always the best model
min_accuracy = val_accuracy # in this way you don't have to care too much about the number of epochs
train_log_loss.close()
train_log_acc.close()
val_log_acc.close()
val_log_loss.close()
