add test scripts

Author: bryanyzhu

main_single_gpu.py

@@ -3,13 +3,15 @@
 import time
 import argparse
 import shutil
+import numpy as np
 
 import torch
 import torch.nn as nn
 import torch.nn.parallel
 import torch.backends.cudnn as cudnn
 import torch.optim
 import torch.utils.data
+import torchvision.models as basic_models
 
 import video_transforms
 import models
@@ -24,28 +26,28 @@
 parser = argparse.ArgumentParser(description='PyTorch Two-Stream Action Recognition')
 parser.add_argument('data', metavar='DIR',
                     help='path to dataset')
-parser.add_argument('--settings', metavar='DIR', default='./settings', 
+parser.add_argument('--settings', metavar='DIR', default='./settings',
                     help='path to datset setting files')
 parser.add_argument('--modality', '-m', metavar='MODALITY', default='rgb',
                     choices=["rgb", "flow"],
                     help='modality: rgb | flow')
 parser.add_argument('--dataset', '-d', default='ucf101',
                     choices=["ucf101", "hmdb51"],
                     help='dataset: ucf101 | hmdb51')
-parser.add_argument('--arch', '-a', metavar='ARCH', default='vgg16',
+parser.add_argument('--arch', '-a', metavar='ARCH', default='rgb_vgg16',
                     choices=model_names,
                     help='model architecture: ' +
                         ' | '.join(model_names) +
-                        ' (default: vgg16)')
+                        ' (default: rgb_vgg16)')
 parser.add_argument('-s', '--split', default=1, type=int, metavar='S',
                     help='which split of data to work on (default: 1)')
 parser.add_argument('-j', '--workers', default=4, type=int, metavar='N',
                     help='number of data loading workers (default: 4)')
-parser.add_argument('--epochs', default=400, type=int, metavar='N',
+parser.add_argument('--epochs', default=250, type=int, metavar='N',
                     help='number of total epochs to run')
 parser.add_argument('--start-epoch', default=0, type=int, metavar='N',
                     help='manual epoch number (useful on restarts)')
-parser.add_argument('-b', '--batch-size', default=50, type=int,
+parser.add_argument('-b', '--batch-size', default=25, type=int,
                     metavar='N', help='mini-batch size (default: 50)')
 parser.add_argument('--iter-size', default=5, type=int,
                     metavar='I', help='iter size as in Caffe to reduce memory usage (default: 5)')
@@ -57,20 +59,24 @@
                     metavar='N', help='resize height (default: 256)')
 parser.add_argument('--lr', '--learning-rate', default=0.001, type=float,
                     metavar='LR', help='initial learning rate')
+parser.add_argument('--lr_steps', default=[100, 200], type=float, nargs="+",
+                    metavar='LRSteps', help='epochs to decay learning rate by 10')
 parser.add_argument('--momentum', default=0.9, type=float, metavar='M',
                     help='momentum')
-parser.add_argument('--weight-decay', '--wd', default=1e-4, type=float,
+parser.add_argument('--weight-decay', '--wd', default=5e-4, type=float,
                     metavar='W', help='weight decay (default: 1e-4)')
-parser.add_argument('--print-freq', default=45, type=int,
+parser.add_argument('--print-freq', default=50, type=int,
                     metavar='N', help='print frequency (default: 20)')
-parser.add_argument('--save-freq', default=40, type=int,
+parser.add_argument('--save-freq', default=25, type=int,
                     metavar='N', help='save frequency (default: 20)')
 parser.add_argument('--resume', default='./checkpoints', type=str, metavar='PATH',
                     help='path to latest checkpoint (default: none)')
 parser.add_argument('-e', '--evaluate', dest='evaluate', action='store_true',
                     help='evaluate model on validation set')
 
 best_prec1 = 0
+os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"   
+os.environ["CUDA_VISIBLE_DEVICES"]="0"
 
 def main():
     global args, best_prec1
@@ -79,7 +85,7 @@ def main():
     # create model
     print("Building model ... ")
     model = build_model()
-    print("Model %s is loaded. " % (args.modality + "_" + args.arch))
+    print("Model %s is loaded. " % (args.arch))
 
     if not os.path.exists(args.resume):
         os.makedirs(args.resume)
@@ -95,26 +101,21 @@ def main():
     cudnn.benchmark = True
 
     # Data transforming
-    # clip_mean = [0.485, 0.456, 0.406] * args.new_length
-    # clip_std = [0.229, 0.224, 0.225] * args.new_length
-    clip_mean = [0.5, 0.5, 0.5] * args.new_length
-    clip_std = [0.5, 0.5, 0.5] * args.new_length
-    normalize = video_transforms.Normalize(mean=clip_mean,
-                                     std=clip_std)
-
     if args.modality == "rgb":
+        is_color = True
         scale_ratios = [1.0, 0.875, 0.75, 0.66]
-    elif args.modality == "flow": 
+        clip_mean = [0.485, 0.456, 0.406] * args.new_length
+        clip_std = [0.229, 0.224, 0.225] * args.new_length
+    elif args.modality == "flow":
+        is_color = False
         scale_ratios = [1.0, 0.875, 0.75]
+        clip_mean = [0.5, 0.5] * args.new_length
+        clip_std = [0.226, 0.226] * args.new_length
     else:
         print("No such modality. Only rgb and flow supported.")
 
-    if args.modality == "rgb":
-        is_color = True
-    elif args.modality == "flow": 
-        is_color = False
-    else:
-        print("No such modality. Only rgb and flow supported.")
+    normalize = video_transforms.Normalize(mean=clip_mean,
+                                     std=clip_std)
 
     train_transform = video_transforms.Compose([
             # video_transforms.Scale((256)),
@@ -130,29 +131,29 @@ def main():
             video_transforms.ToTensor(),
             normalize,
         ])
-    
-    # data loading 
+
+    # data loading
     train_setting_file = "train_%s_split%d.txt" % (args.modality, args.split)
     train_split_file = os.path.join(args.settings, args.dataset, train_setting_file)
     val_setting_file = "val_%s_split%d.txt" % (args.modality, args.split)
     val_split_file = os.path.join(args.settings, args.dataset, val_setting_file)
     if not os.path.exists(train_split_file) or not os.path.exists(val_split_file):
         print("No split file exists in %s directory. Preprocess the dataset first" % (args.settings))
 
-    train_dataset = datasets.__dict__[args.dataset](root=args.data, 
-                                                    source=train_split_file, 
-                                                    phase="train", 
+    train_dataset = datasets.__dict__[args.dataset](root=args.data,
+                                                    source=train_split_file,
+                                                    phase="train",
                                                     modality=args.modality,
-                                                    is_color=is_color, 
+                                                    is_color=is_color,
                                                     new_length=args.new_length,
                                                     new_width=args.new_width,
                                                     new_height=args.new_height,
                                                     video_transform=train_transform)
-    val_dataset = datasets.__dict__[args.dataset](root=args.data, 
-                                                  source=val_split_file, 
-                                                  phase="val", 
-                                                  modality=args.modality, 
-                                                  is_color=is_color, 
+    val_dataset = datasets.__dict__[args.dataset](root=args.data,
+                                                  source=val_split_file,
+                                                  phase="val",
+                                                  modality=args.modality,
+                                                  is_color=is_color,
                                                   new_length=args.new_length,
                                                   new_width=args.new_width,
                                                   new_height=args.new_height,
@@ -182,7 +183,9 @@ def main():
         train(train_loader, model, criterion, optimizer, epoch)
 
         # evaluate on validation set
-        prec1 = validate(val_loader, model, criterion)
+        prec1 = 0.0
+        if (epoch + 1) % args.save_freq == 0:
+            prec1 = validate(val_loader, model, criterion)
 
         # remember best prec@1 and save checkpoint
         is_best = prec1 > best_prec1
@@ -200,8 +203,7 @@ def main():
 
 def build_model():
 
-    model_name = args.modality + "_" + args.arch
-    model = models.__dict__[model_name](pretrained=True, num_classes=101)
+    model = models.__dict__[args.arch](pretrained=True, num_classes=101)
     model.cuda()
     return model
 
@@ -256,7 +258,6 @@ def train(train_loader, model, criterion, optimizer, epoch):
                       'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
                       'Prec@1 {top1.val:.3f} ({top1.avg:.3f})'.format(
                        epoch, i+1, len(train_loader)+1, batch_time=batch_time, loss=losses, top1=top1))
-            
 
 def validate(val_loader, model, criterion):
     batch_time = AverageMeter()
@@ -302,15 +303,13 @@ def validate(val_loader, model, criterion):
 
     return top1.avg
 
-
 def save_checkpoint(state, is_best, filename, resume_path):
     cur_path = os.path.join(resume_path, filename)
     best_path = os.path.join(resume_path, 'model_best.pth.tar')
     torch.save(state, cur_path)
     if is_best:
         shutil.copyfile(cur_path, best_path)
 
-
 class AverageMeter(object):
     """Computes and stores the average and current value"""
     def __init__(self):
@@ -328,14 +327,13 @@ def update(self, val, n=1):
         self.count += n
         self.avg = self.sum / self.count
 
-
 def adjust_learning_rate(optimizer, epoch):
     """Sets the learning rate to the initial LR decayed by 10 every 150 epochs"""
-    lr = args.lr * (0.1 ** (epoch // 150))
-    print(lr)
+    decay = 0.1 ** (sum(epoch >= np.array(args.lr_steps)))
+    lr = args.lr * decay
+    print("Current learning rate is %4.6f:" % lr)
     for param_group in optimizer.param_groups:
         param_group['lr'] = lr
-        # param_group['lr'] = param_group['lr']/2
 
 def accuracy(output, target, topk=(1,)):
     """Computes the precision@k for the specified values of k"""

models/__init__.py

@@ -1,2 +1,4 @@
 from .rgb_vgg16 import *
 from .flow_vgg16 import *
+from .rgb_resnet import *
+from .flow_resnet import *

models/flow_vgg16.py

@@ -101,7 +101,8 @@ def flow_vgg16(pretrained=False, **kwargs):
         pretrained (bool): If True, returns a model pre-trained on ImageNet
     """
     model = VGG(make_layers(cfg['D']), **kwargs)
-    in_channels = 20
+    # TODO: hardcoded for now for 10 optical flow images, set it as an argument later 
+    in_channels = 20            
     if pretrained:
         # model.load_state_dict(model_zoo.load_url(model_urls['vgg16']))
         pretrained_dict = model_zoo.load_url(model_urls['vgg16'])

scripts/eval_ucf101_pytorch/VideoSpatialPrediction.py

@@ -0,0 +1,102 @@
+'''
+A sample function for classification using spatial network
+Customize as needed:
+e.g. num_categories, layer for feature extraction, batch_size
+'''
+
+import os
+import sys
+import numpy as np
+import math
+import cv2
+import scipy.io as sio
+
+import torch
+import torch.nn as nn
+import torch.nn.parallel
+import torch.backends.cudnn as cudnn
+import torch.optim
+import torch.utils.data
+import torchvision.transforms as transforms
+import torchvision.datasets as datasets
+import torchvision.models as models
+
+sys.path.insert(0, "../../")
+import video_transforms
+
+def VideoSpatialPrediction(
+        vid_name,
+        net,
+        num_categories,
+        start_frame=0,
+        num_frames=0,
+        num_samples=25
+        ):
+
+    if num_frames == 0:
+        imglist = os.listdir(vid_name)
+        duration = len(imglist)
+        # print(duration)
+    else:
+        duration = num_frames
+
+    clip_mean = [0.485, 0.456, 0.406]
+    clip_std = [0.229, 0.224, 0.225]
+    normalize = video_transforms.Normalize(mean=clip_mean,
+                                     std=clip_std)
+    val_transform = video_transforms.Compose([
+            video_transforms.ToTensor(),
+            normalize,
+        ])
+
+    # selection
+    step = int(math.floor((duration-1)/(num_samples-1)))
+    dims = (256,340,3,num_samples)
+    rgb = np.zeros(shape=dims, dtype=np.float64)
+    rgb_flip = np.zeros(shape=dims, dtype=np.float64)
+
+    for i in range(num_samples):
+        img_file = os.path.join(vid_name, 'image_{0:04d}.jpg'.format(i*step+1))
+        img = cv2.imread(img_file, cv2.IMREAD_UNCHANGED)
+        img = cv2.resize(img, dims[1::-1])
+        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+        rgb[:,:,:,i] = img
+        rgb_flip[:,:,:,i] = img[:,::-1,:]
+
+    # crop
+    rgb_1 = rgb[:224, :224, :,:]
+    rgb_2 = rgb[:224, -224:, :,:]
+    rgb_3 = rgb[16:240, 60:284, :,:]
+    rgb_4 = rgb[-224:, :224, :,:]
+    rgb_5 = rgb[-224:, -224:, :,:]
+    rgb_f_1 = rgb_flip[:224, :224, :,:]
+    rgb_f_2 = rgb_flip[:224, -224:, :,:]
+    rgb_f_3 = rgb_flip[16:240, 60:284, :,:]
+    rgb_f_4 = rgb_flip[-224:, :224, :,:]
+    rgb_f_5 = rgb_flip[-224:, -224:, :,:]
+
+    rgb = np.concatenate((rgb_1,rgb_2,rgb_3,rgb_4,rgb_5,rgb_f_1,rgb_f_2,rgb_f_3,rgb_f_4,rgb_f_5), axis=3)
+
+    _, _, _, c = rgb.shape
+    rgb_list = []
+    for c_index in range(c):
+        cur_img = rgb[:,:,:,c_index].squeeze()
+        cur_img_tensor = val_transform(cur_img)
+        rgb_list.append(np.expand_dims(cur_img_tensor.numpy(), 0))
+        
+    rgb_np = np.concatenate(rgb_list,axis=0)
+    # print(rgb_np.shape)
+    batch_size = 25
+    prediction = np.zeros((num_categories,rgb.shape[3]))
+    num_batches = int(math.ceil(float(rgb.shape[3])/batch_size))
+
+    for bb in range(num_batches):
+        span = range(batch_size*bb, min(rgb.shape[3],batch_size*(bb+1)))
+        input_data = rgb_np[span,:,:,:]
+        imgDataTensor = torch.from_numpy(input_data).type(torch.FloatTensor).cuda()
+        imgDataVar = torch.autograd.Variable(imgDataTensor)
+        output = net(imgDataVar)
+        result = output.data.cpu().numpy()
+        prediction[:, span] = np.transpose(result)
+
+    return prediction