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]