rootDir=rootDir,
N_FRAME=N_FRAME) # for raw RGB frames
# trainingData = videoDatasetPreGenOF(folderList=trainFoldeList, rootDir=rootDir, N_FRAME=N_FRAME) # for pre-generated OF
dataloader = DataLoader(trainingData,
batch_size=BATCH_SIZE,
shuffle=True,
num_workers=1)
## Initializing r, theta
stateDict = torch.load(trained_encoder)['state_dict']
Dtheta = stateDict['l1.theta']
Drr = stateDict['l1.rr']
## Create the model
model = unfrozen_DyanC(Drr, Dtheta, T, PRE, gpu_id)
# model = DyanC(T, PRE, pretrained_dyan, gpu_id);
model.cuda(gpu_id)
model.train()
stateDict = torch.load(trained_encoder)['state_dict']
Dtheta = stateDict['l1.theta']
Drr = stateDict['l1.rr']
baseDyan = OFModel(Drr, Dtheta, T, PRE, gpu_id)
baseDyan.eval()
optimizer = torch.optim.Adam(model.parameters(), lr=LR)
scheduler = lr_scheduler.MultiStepLR(
optimizer, milestones=[50,
100], gamma=0.1) # if Kitti: milestones=[100,150]
baseDyan = OFModel(Drr, Dtheta, T, PRE, gpu_id)
baseDyan.cuda(gpu_id)
baseDyan.eval()
'''
loadedcheckpoint = torch.load(opticalflow_ckpt_file)
stateDict = loadedcheckpoint['state_dict']
# load parameters
Dtheta = stateDict['l1.theta']
Drr = stateDict['l1.rr']
dyanEncoder = OFModel(Drr, Dtheta, FRA,PRE,gpu_id)
dyanEncoder.cuda(gpu_id)
'''
## Load the classifier network
classifier = unfrozen_DyanC(Drr, Dtheta, T, PRE, gpu_id)
classifier.load_state_dict(torch.load(classifier_ckpt_file)["state_dict"])
classifier.cuda(gpu_id)
classifier.eval()
ofSample = torch.FloatTensor(2, FRA, numOfPixels)
# set test list name:
validateFolderFile = 'validatelist01.txt'
# set test data directory:
# rootDir = '/storage/truppr/UCF-FLOWS-FULL/'
rootDir = '/data/Abhishek/frames/'
# for UCF dataset:
validateFoldeList = getListOfFolders(validateFolderFile)[::5]
dataloader = DataLoader(trainingData,
batch_size=BATCH_SIZE,
shuffle=True,
num_workers=1)
tvnet_args = arguments().parse() # tvnet parameters
tvnet_args.batch_size = N_FRAME
tvnet_args.data_size = [tvnet_args.batch_size, 3, 240, 320]
## Initializing r, theta
stateDict = torch.load(trained_encoder)['state_dict']
Dtheta = stateDict['l1.theta']
Drr = stateDict['l1.rr']
## Create the model
model = unfrozen_DyanC(Drr, Dtheta, T, PRE, tvnet_args, gpu_id)
# model = DyanC(T, PRE, pretrained_dyan, gpu_id);
model.cuda(gpu_id)
model.train()
'''
stateDict = torch.load(pretrained_dyan)['state_dict']
Dtheta = stateDict['l1.theta']
Drr = stateDict['l1.rr']
baseDyan = OFModel(Drr, Dtheta, T, PRE, gpu_id)
baseDyan.eval()
'''
optimizer = torch.optim.Adam(model.parameters(), lr=LR)
scheduler = lr_scheduler.MultiStepLR(
optimizer, milestones=[50,
100], gamma=0.1) # if Kitti: milestones=[100,150]
