update readme

README.md

@@ -1,10 +1,46 @@
 # ReID_baseline
 Baseline model (with bottleneck) for person ReID (using softmax and triplet loss).
 
-## Learning rate
-This is the warmpup strategy learning rate
+We support
+- multi-GPU training
+- easy dataset preparation
+- end-to-end training and evaluation
 
-<img src='https://ws3.sinaimg.cn/large/006tNc79ly1fthmcjwoaaj31kw0natad.jpg' height='200'>
+## Get Started
+1. `cd` to folder where you want to download this repo
+2. Run `git clone https://github.com/L1aoXingyu/reid_baseline.git`
+3. Install dependencies:
+    - [pytorch](https://pytorch.org/)
+    - torchvision
+    - tensorflow (for tensorboard)
+    - [tensorboardX](https://github.com/lanpa/tensorboardX)
+4. Prepare dataset
+    
+    Create a directory to store reid datasets under this repo via
+    ```bash
+    cd reid_baseline
+    mkdir data
+    ```
+    1. Download dataset to `data/` from http://www.liangzheng.org/Project/project_reid.html
+    2. Extract dataset and rename to `market1501`. The data structure would like:
+    ```
+    market1501/
+        bounding_box_test/
+        bounding_box_train/
+    ```
+5. Prepare pretrained model if you don't have
+    ```python
+    from torchvision import models
+    models.resnet50(pretrained=True)
+    ```
+    Then it will automatically download model in `~.torch/models/`, you should set this path in `config.py`
+
+## Train
+You can run 
+```bash
+bash scripts/train_triplet_softmax.sh
+```
+in `reid_baseline` folder if you want to train with softmax and triplet loss. You can find others train scripts in `scripts`.
 
 ## Results
 

bases/base_trainer.py

@@ -60,7 +60,7 @@ class BaseTrainer(object):
                               data_time.val, data_time.mean,
                               losses.val, losses.mean))
         param_group = self.optimizer.param_groups
-        print('Epoch: [{}]\tEpoch Time {:.3f} s\tLoss {:.3e}\t'
+        print('Epoch: [{}]\tEpoch Time {:.3f} s\tLoss {:.3f}\t'
               'Lr {:.2e}'
               .format(epoch, batch_time.sum, losses.mean, param_group[0]['lr']))
         print()

config.py

@@ -37,6 +37,7 @@ class DefaultConfig(object):
     # model options
     model_name = 'softmax'  # softmax, triplet, softmax_triplet
     last_stride = 1
+    pretrained_model = '/home/test2/.torch/models/resnet50-19c8e357.pth'
 
     # miscs
     print_freq = 30

datasets/__init__.py

@@ -0,0 +1,11 @@
+# encoding: utf-8
+"""
+@author:  liaoxingyu
+@contact: sherlockliao01@gmail.com
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+from __future__ import unicode_literals
+

datasets/data_manager.py

@@ -20,7 +20,7 @@ class Market1501(object):
     """
     dataset_dir = 'market1501'
 
-    def __init__(self, root='/home/liaoxingyu/', **kwargs):
+    def __init__(self, root='data', **kwargs):
         self.dataset_dir = osp.join(root, self.dataset_dir)
         self.train_dir = osp.join(self.dataset_dir, 'bounding_box_train')
         self.query_dir = osp.join(self.dataset_dir, 'query')

models/networks.py

@@ -12,7 +12,7 @@ from __future__ import unicode_literals
 from .baseline_model import ResNetBuilder
 
 
-def get_baseline_model(num_classes, last_stride=1, model_path='/DATA/model_zoo/resnet50-19c8e357.pth'):
+def get_baseline_model(num_classes, last_stride=1, model_path='/home/test2/.torch/models/resnet50-19c8e357.pth'):
     model = ResNetBuilder(num_classes, last_stride, model_path)
     optim_policy = model.get_optim_policy()
     return model, optim_policy

scripts/train_softmax.sh

@@ -1,4 +1,4 @@
 #!/usr/bin/env bash
 
-CUDA_VISIBLE_DEVICES=0 python3 main_reid.py train --save_dir='/DATA/pytorch-ckpt/market1501_softmax' --max_epoch=400 \
+CUDA_VISIBLE_DEVICES=0 python3 main_reid.py train --save_dir='./pytorch-ckpt/market1501_softmax' --max_epoch=400 \
 --eval_step=50 --model_name='softmax' 

scripts/train_triplet.sh

@@ -1,4 +1,4 @@
 #!/usr/bin/env bash
 
-CUDA_VISIBLE_DEVICES=0 python3 main_reid.py train --save_dir='/DATA/pytorch-ckpt/market1501_triplet' --max_epoch=400 \
+CUDA_VISIBLE_DEVICES=0 python3 main_reid.py train --save_dir='./pytorch-ckpt/market1501_triplet' --max_epoch=400 \
 --eval_step=50 --model_name='triplet' 

scripts/train_triplet_softmax.sh

@@ -1,4 +1,4 @@
 #!/usr/bin/env bash
 
-CUDA_VISIBLE_DEVICES=0 python3 main_reid.py train --save_dir='/DATA/pytorch-ckpt/market1501_softmax_triplet' \
---max_epoch=400 --eval_step=50 --model_name='softmax_triplet'
+CUDA_VISIBLE_DEVICES=0 python3 main_reid.py train --save_dir='./pytorch-ckpt/market1501_softmax_triplet' \
+--max_epoch=400 --eval_step=5 --model_name='softmax_triplet' \