update mscoco support. add annotations json file (see readme)

Author: bearpaw

.gitignore

@@ -11,7 +11,8 @@ test_transforms.py
 experiments
 run
 data/mscoco/coco
-data/mscoco/keypoint
+data/mscoco/annotations
+data/mscoco/images
 *.m~
 miscs/posetrack
 miscs/h36m

README.md

@@ -4,7 +4,7 @@ PyTorch-Pose is a PyTorch implementation of the general pipeline for 2D single h
 
 Some codes for data preparation and augmentation are brought from the [Stacked hourglass network](https://github.com/anewell/pose-hg-train). Thanks to the original author.
 
-**Update: this repository is compatible with PyTorch 1.0 now!**
+**Update: this repository is compatible with PyTorch 0.4.1/1.0 now!**
 
 ## Features
 - Multi-thread data loading
@@ -13,7 +13,7 @@ Some codes for data preparation and augmentation are brought from the [Stacked h
 - Training/testing results visualization
 
 ## Installation
-1. PyTorch (>= 0.4.0): Please follow the [installation instruction of PyTorch](http://pytorch.org/). Note that the code is developed with Python2 and has not been tested with Python3 yet.
+1. PyTorch (>= 0.4.1): Please follow the [installation instruction of PyTorch](http://pytorch.org/). Note that the code is developed with Python2 and has not been tested with Python3 yet.
 
 2. Clone the repository with submodule
    ```
@@ -24,12 +24,17 @@ Some codes for data preparation and augmentation are brought from the [Stacked h
    ```
    ln -s PATH_TO_MPII_IMAGES_DIR data/mpii/images
    ```
+   For training/testing on COCO, please refer to [COCO Readme](https://github.com/bearpaw/pytorch-pose/blob/master/data/mscoco/README.md).
 
 4. Modify your `.bashrc` file:
    ```
    export PYTHONPATH=".:$PYTHONPATH"
    ```
 
+5. Download annotation file:
+    * (MPII) Download [mpii_annotations.json](https://drive.google.com/open?id=1mQrH_yVHeB93rzCfyq5kC9ZYTwZeMsMm) and save it to `data/mpii`
+    * (MSCOCO) Download [coco_annotations_2014.json](https://drive.google.com/open?id=1jrxis4ujrLlkwoD2GOdv3PGzygpQ04k7) or/and [coco_annotations_2017.json](https://drive.google.com/open?id=1YuzpScAfzemwZqUuZBrbBZdoplXEqUse) and save it to `data/mscoco`
+
 ## Usage
 
 ### Testing

data/lsp/images/lsp_dataset

@@ -1,9 +1,22 @@
 ## Directory structure
+To train on coco (single person pose estimation), place the following files in the right place.
 
-- `coco`: [coco API](https://github.com/pdollar/coco)
-- `keypoint`: COCO keypoint dataset
-    - `images`: tain and val datasets
-        - `train2014`
-        - `val2014`
-    - `person_keypoints_train+val5k2014`: annotations (JSON files)
-- `coco_annotations.json`: reformatted annotation generated by `./miscs/gen_coco.m`
+- `cocoapi`: [coco API](https://github.com/cocodataset/cocoapi)
+- `annotations`: COCO keypoint dataset
+  - `person_keypoints_train+val5k2014`: annotations (JSON files)
+- `images`: tain and val datasets
+    - `train2014`
+    - `val2014`
+- `coco_annotations_2014.json` or `coco_annotations_2014.json`: reformatted annotation generated by `./miscs/gen_coco.m`
+
+### Prepare `json` annotation file
+Please edit `year = 2017; % 2014 or 2017 ` in [gen_coco.m](https://github.com/bearpaw/pytorch-pose/blob/master/miscs/gen_coco.m) to generate the corresponding annotation file (coco2014 or coco2017).
+
+### Training recipes
+
+```
+CUDA_VISIBLE_DEVICES=1 python example/mscoco.py \
+--checkpoint checkpoint/mscoco/hg-s1-b1 \
+-s 1 \
+-b 1
+```

data/lsp/images/lspet_dataset

@@ -46,7 +46,7 @@ def main(args):
     # define loss function (criterion) and optimizer
     criterion = torch.nn.MSELoss(size_average=True).cuda()
 
-    optimizer = torch.optim.RMSprop(model.parameters(), 
+    optimizer = torch.optim.RMSprop(model.parameters(),
                                 lr=args.lr,
                                 momentum=args.momentum,
                                 weight_decay=args.weight_decay)
@@ -66,7 +66,7 @@ def main(args):
             logger = Logger(join(args.checkpoint, 'log.txt'), title=title, resume=True)
         else:
             print("=> no checkpoint found at '{}'".format(args.resume))
-    else:        
+    else:
         logger = Logger(join(args.checkpoint, 'log.txt'), title=title)
         logger.set_names(['Epoch', 'LR', 'Train Loss', 'Val Loss', 'Train Acc', 'Val Acc'])
 
@@ -79,15 +79,15 @@ def main(args):
                       sigma=args.sigma, label_type=args.label_type),
         batch_size=args.train_batch, shuffle=True,
         num_workers=args.workers, pin_memory=True)
-    
+
     val_loader = torch.utils.data.DataLoader(
         datasets.LSP('data/lsp/LEEDS_annotations.json', 'data/lsp/images',
                       sigma=args.sigma, label_type=args.label_type, train=False),
         batch_size=args.test_batch, shuffle=False,
         num_workers=args.workers, pin_memory=True)
 
     if args.evaluate:
-        print('\nEvaluation only') 
+        print('\nEvaluation only')
         loss, acc, predictions = validate(val_loader, model, criterion, args.num_classes, args.debug, args.flip)
         save_pred(predictions, checkpoint=args.checkpoint)
         return
@@ -174,7 +174,7 @@ def train(train_loader, model, criterion, optimizer, debug=False, flip=True):
             plt.draw()
 
         # measure accuracy and record loss
-        losses.update(loss.data[0], inputs.size(0))
+        losses.update(loss.item(), inputs.size(0))
         acces.update(acc[0], inputs.size(0))
 
         # compute gradient and do SGD step
@@ -232,7 +232,7 @@ def validate(val_loader, model, criterion, num_classes, debug=False, flip=True):
         score_map = output[-1].data.cpu()
         if flip:
             flip_input_var = torch.autograd.Variable(
-                    torch.from_numpy(fliplr(inputs.clone().numpy())).float().cuda(), 
+                    torch.from_numpy(fliplr(inputs.clone().numpy())).float().cuda(),
                     volatile=True
                 )
             flip_output_var = model(flip_input_var)
@@ -267,7 +267,7 @@ def validate(val_loader, model, criterion, num_classes, debug=False, flip=True):
             plt.draw()
 
         # measure accuracy and record loss
-        losses.update(loss.data[0], inputs.size(0))
+        losses.update(loss.item(), inputs.size(0))
         acces.update(acc[0], inputs.size(0))
 
         # measure elapsed time
@@ -348,4 +348,4 @@ def validate(val_loader, model, criterion, num_classes, debug=False, flip=True):
                         help='show intermediate results')
 
 
-    main(parser.parse_args())
+    main(parser.parse_args())

data/mscoco/README.md

@@ -124,7 +124,7 @@ def main(args):
             'state_dict': model.state_dict(),
             'best_acc': best_acc,
             'optimizer' : optimizer.state_dict(),
-        }, predictions, is_best, checkpoint=args.checkpoint)
+        }, predictions, is_best, checkpoint=args.checkpoint, snapshot=args.snapshot)
 
     logger.close()
     logger.plot(['Train Acc', 'Val Acc'])
@@ -229,13 +229,10 @@ def validate(val_loader, model, criterion, num_classes, debug=False, flip=True):
         target = target.to(device, non_blocking=True)
 
         # compute output
-        output = model(input)
+        output = model(input)[-1]
         score_map = output[-1].cpu()
         if flip:
-            flip_input = torch.autograd.Variable(
-                    torch.from_numpy(fliplr(input.clone().numpy())).float().to(device),
-                    volatile=True
-                )
+            flip_input = torch.from_numpy(fliplr(input.clone().numpy())).float().to(device)
             flip_output_var = model(flip_input)
             flip_output = flip_back(flip_output_var[-1].cpu())
             score_map += flip_output