20 lines
694 B
Python
Raw Normal View History

2019-03-19 17:26:08 +00:00
import torchreid
2019-03-21 12:53:21 +00:00
datamanager = torchreid.data.ImageDataManager(
2019-03-19 17:26:08 +00:00
root='data',
2019-03-21 12:53:21 +00:00
sources=['market1501', 'cuhk03'],
targets='market1501',
2019-03-19 17:26:08 +00:00
height=128,
2019-03-21 12:53:21 +00:00
width=64,
combineall=True,
batch_size=5
2019-03-19 17:26:08 +00:00
)
model = torchreid.models.build_model(
name='squeezenet1_0',
2019-03-21 12:53:21 +00:00
num_classes=datamanager.num_train_pids,
2019-03-19 17:26:08 +00:00
loss='softmax'
)
optimizer = torchreid.optim.build_optimizer(model)
scheduler = torchreid.optim.build_lr_scheduler(optimizer, lr_scheduler='multi_step', stepsize=[10, 20])
2019-03-21 12:53:21 +00:00
engine = torchreid.engine.ImageSoftmaxEngine(datamanager, model, optimizer, scheduler=scheduler)
#engine.run(max_epoch=1, print_freq=1, fixbase_epoch=0, open_layers='classifier', test_only=True)