def main():
# Model Initialize
m = createModel().cuda()
if opt.loadModel:
print('Loading Model from {}'.format(opt.loadModel))
m.load_state_dict(torch.load(opt.loadModel))
if not os.path.exists("../exp/{}/{}".format(opt.dataset, opt.expID)):
try:
os.mkdir("../exp/{}/{}".format(opt.dataset, opt.expID))
except FileNotFoundError:
os.mkdir("../exp/{}".format(opt.dataset))
os.mkdir("../exp/{}/{}".format(opt.dataset, opt.expID))
else:
print('Create new model')
if not os.path.exists("../exp/{}/{}".format(opt.dataset, opt.expID)):
try:
os.mkdir("../exp/{}/{}".format(opt.dataset, opt.expID))
except FileNotFoundError:
os.mkdir("../exp/{}".format(opt.dataset))
os.mkdir("../exp/{}/{}".format(opt.dataset, opt.expID))
criterion = torch.nn.MSELoss().cuda()
if opt.optMethod == 'rmsprop':
optimizer = torch.optim.RMSprop(m.parameters(),
lr=opt.LR,
momentum=opt.momentum,
weight_decay=opt.weightDecay)
elif opt.optMethod == 'adam':
optimizer = torch.optim.Adam(m.parameters(), lr=opt.LR)
else:
raise Exception
writer = SummaryWriter('.tensorboard/{}/{}'.format(opt.dataset, opt.expID))
# Prepare Dataset
if opt.dataset == 'coco':
train_dataset = coco.Mscoco(train=True)
val_dataset = coco.Mscoco(train=False)
train_dataset = coco.Mscoco(train=True)
val_dataset = coco.Mscoco(train=False)
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=opt.trainBatch,
shuffle=True,
num_workers=opt.nThreads,
pin_memory=True)
val_loader = torch.utils.data.DataLoader(val_dataset,
batch_size=opt.validBatch,
shuffle=False,
num_workers=opt.nThreads,
pin_memory=True)
# Model Transfer
m = torch.nn.DataParallel(m).cuda()
# Start Training
for i in range(opt.nEpochs):
opt.epoch = i
print('############# Starting Epoch {} #############'.format(
opt.epoch))
loss, acc = train(train_loader, m, criterion, optimizer, writer)
print('Train-{idx:d} epoch | loss:{loss:.8f} | acc:{acc:.4f}'.format(
idx=opt.epoch, loss=loss, acc=acc))
opt.acc = acc
opt.loss = loss
m_dev = m.module
if i % opt.snapshot == 0:
torch.save(
m_dev.state_dict(),
'../exp/{}/{}/model_{}.pkl'.format(opt.dataset, opt.expID,
opt.epoch))
torch.save(
opt, '../exp/{}/{}/option.pkl'.format(opt.dataset, opt.expID,
opt.epoch))
torch.save(
optimizer,
'../exp/{}/{}/optimizer.pkl'.format(opt.dataset, opt.expID))
loss, acc = valid(val_loader, m, criterion, optimizer, writer)
print('Valid-{idx:d} epoch | loss:{loss:.8f} | acc:{acc:.4f}'.format(
idx=i, loss=loss, acc=acc))
'''
if opt.dataset != 'mpii':
with torch.no_grad():
mAP, mAP5 = prediction(m)
print('Prediction-{idx:d} epoch | mAP:{mAP:.3f} | mAP0.5:{mAP5:.3f}'.format(
idx=i,
mAP=mAP,
mAP5=mAP5
))
'''
writer.close()
def main():
print(opt)
# Model Initialize
m = createModel().cuda()
if opt.loadModel:
print('Loading Model from {}'.format(opt.loadModel))
'''
ckp = torch.load(opt.loadModel)
for name,param in m.state_dict().items():
if name in ckp:
ckp_param = ckp[name]
if ckp_param.shape == param.shape:
param.copy_(ckp_param)
print(name, 'copy successfully')
else:
print(name, 'shape is inconsistent with checkpoint')
else:
print(name, 'can not be found in checkpoint')
'''
m.load_state_dict(torch.load(opt.loadModel))
if not os.path.exists("../exp/{}/{}".format(opt.dataset, opt.expID)):
try:
os.mkdir("../exp/{}/{}".format(opt.dataset, opt.expID))
except FileNotFoundError:
os.mkdir("../exp/{}".format(opt.dataset))
os.mkdir("../exp/{}/{}".format(opt.dataset, opt.expID))
else:
print('Create new model')
# import pdb;pdb.set_trace()
if not os.path.exists("../exp/{}/{}".format(opt.dataset, opt.expID)):
try:
os.mkdir("../exp/{}/{}".format(opt.dataset, opt.expID))
except FileNotFoundError:
os.mkdir("../exp/{}".format(opt.dataset))
os.mkdir("../exp/{}/{}".format(opt.dataset, opt.expID))
criterion = torch.nn.MSELoss().cuda()
if opt.optMethod == 'rmsprop':
optimizer = torch.optim.RMSprop(m.parameters(),
lr=opt.LR,
momentum=opt.momentum,
weight_decay=opt.weightDecay)
elif opt.optMethod == 'adam':
optimizer = torch.optim.Adam(m.parameters(), lr=opt.LR)
elif opt.optMethod == 'rmsprop_refine':
print('opt rmsprop refine')
optimizer = torch.optim.RMSprop([{
"params": m.preact.parameters(),
"lr": 1e-5
}, {
"params": m.suffle1.parameters(),
"lr": 1e-5
}, {
"params": m.duc1.parameters(),
"lr": 1e-5
}, {
"params": m.duc2.parameters(),
"lr": 1e-5
}, {
"params": m.conv_out.parameters(),
"lr": 1e-4
}],
lr=opt.LR,
momentum=opt.momentum,
weight_decay=opt.weightDecay)
else:
raise Exception
if opt.loadOptimizer:
optimizer = torch.load(opt.loadOptimizer)
writer = SummaryWriter('.tensorboard/{}/{}'.format(opt.dataset, opt.expID))
# Prepare Dataset
if opt.dataset == 'coco':
train_dataset = coco.Mscoco(train=True)
val_dataset = coco.Mscoco(train=False)
elif opt.dataset == 'h36m':
train_dataset = h36m.H36M(train=True)
val_dataset = h36m.H36M(train=False)
'''
train_loader = torch.utils.data.DataLoader(
train_dataset, batch_size=opt.trainBatch, shuffle=False, sampler=RandomSampler(train_dataset, replacement=True, num_samples=len(train_dataset)//10), num_workers=opt.nThreads, pin_memory=True)
val_loader = torch.utils.data.DataLoader(
val_dataset, batch_size=opt.validBatch, shuffle=False, sampler=RandomSampler(val_dataset, replacement=True, num_samples=len(val_dataset)//10), num_workers=opt.nThreads, pin_memory=True)
'''
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=opt.trainBatch,
shuffle=True,
num_workers=opt.nThreads,
pin_memory=True)
val_loader = torch.utils.data.DataLoader(
val_dataset,
batch_size=opt.validBatch,
shuffle=False,
sampler=RandomSampler(val_dataset,
replacement=True,
num_samples=len(val_dataset) // 8),
num_workers=opt.nThreads,
pin_memory=True)
#show_loader_image('train_check_images', train_loader, joint_names=train_dataset.joint_names)
#show_loader_image('valid_check_images', val_loader, joint_names=val_dataset.joint_names)
#return
# Model Transfer
m = torch.nn.DataParallel(m).cuda()
# Start Training
for i in range(opt.nEpochs):
opt.epoch = i
print('############# Starting Epoch {} #############'.format(
opt.epoch))
loss, acc = train(train_loader, m, criterion, optimizer, writer)
print('Train-{idx:d} epoch | loss:{loss:.8f} | acc:{acc:.4f}'.format(
idx=opt.epoch, loss=loss, acc=acc))
opt.acc = acc
opt.loss = loss
m_dev = m.module
if i % opt.snapshot == 0:
torch.save(
m_dev.state_dict(),
'../exp/{}/{}/model_{}.pkl'.format(opt.dataset, opt.expID,
opt.epoch))
torch.save(
opt, '../exp/{}/{}/option.pkl'.format(opt.dataset, opt.expID,
opt.epoch))
torch.save(
optimizer,
'../exp/{}/{}/optimizer.pkl'.format(opt.dataset, opt.expID))
loss, acc = valid(val_loader, m, criterion, optimizer, writer)
print('Valid-{idx:d} epoch | loss:{loss:.8f} | acc:{acc:.4f}'.format(
idx=i, loss=loss, acc=acc))
'''
if opt.dataset != 'mpii':
with torch.no_grad():
mAP, mAP5 = prediction(m)
print('Prediction-{idx:d} epoch | mAP:{mAP:.3f} | mAP0.5:{mAP5:.3f}'.format(
idx=i,
mAP=mAP,
mAP5=mAP5
))
'''
writer.close()
def main():
# Model Initialize
m = createModel().cuda()
m.load_state_dict(torch.load('duc_se.pth'))
criterion = torch.nn.MSELoss().cuda()
if opt.optMethod == 'rmsprop':
optimizer = torch.optim.RMSprop(m.parameters(),
lr=opt.LR,
momentum=opt.momentum,
weight_decay=opt.weightDecay)
elif opt.optMethod == 'adam':
optimizer = torch.optim.Adam(m.parameters(), lr=opt.LR)
else:
raise Exception
writer = SummaryWriter('.tensorboard/{}/{}'.format(opt.dataset, opt.expID))
# Prepare Dataset
if opt.dataset == 'coco':
train_dataset = coco.Mscoco(train=True)
val_dataset = coco.Mscoco(train=False)
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=opt.trainBatch,
shuffle=True,
num_workers=opt.nThreads,
pin_memory=True)
val_loader = torch.utils.data.DataLoader(val_dataset,
batch_size=opt.validBatch,
shuffle=False,
num_workers=opt.nThreads,
pin_memory=True)
# Model Transfer
m = torch.nn.DataParallel(m).cuda()
# Start Training
for i in range(opt.nEpochs):
opt.epoch = i
print('############# Starting Epoch {} #############'.format(
opt.epoch))
loss, acc = train(train_loader, m, criterion, optimizer, writer)
print('Train-{idx:d} epoch | loss:{loss:.8f} | acc:{acc:.4f}'.format(
idx=opt.epoch, loss=loss, acc=acc))
opt.acc = acc
opt.loss = loss
m_dev = m.module
if i % opt.snapshot == 0:
torch.save(
m_dev.state_dict(),
'../exp/{}/{}/model_{}.pkl'.format(opt.dataset, opt.expID,
opt.epoch))
torch.save(
opt, '../exp/{}/{}/option.pkl'.format(opt.dataset, opt.expID,
opt.epoch))
torch.save(
optimizer,
'../exp/{}/{}/optimizer.pkl'.format(opt.dataset, opt.expID))
loss, acc = valid(val_loader, m, criterion, optimizer, writer)
print('Valid-{idx:d} epoch | loss:{loss:.8f} | acc:{acc:.4f}'.format(
idx=i, loss=loss, acc=acc))
writer.close()
def main():
# 初始化模型
m = createModel().cuda()
# 加载已经训练道模型接着训练
if opt.loadModel:
print("加载已训练道模型".format(opt.loadModel))
m.load_state_dict(torch.load(opt.loadModel))
pass
# 新建模型进行训练
else:
print('新建一个模型')
if not os.path.exists("../exp/{}/{}".format(opt.dataset, opt.expID)):
try:
os.mkdir("../exp/{}/{}".format(opt.dataset, opt.expID))
except FileNotFoundError:
os.mkdir("../exp/{}".format(opt.dataset))
os.mkdir("../exp/{}/{}".format(opt.dataset, opt.expID))
criterion = torch.nn.MSELoss().cuda()
# 定义优化器
if opt.optMethod == 'adam':
optimizer = torch.optim.Adam(
m.parameters(),
lr=opt.LR
)
elif opt.optMethod == 'rmsprop':
optimizer = torch.optim.RMSprop(m.parameters(),
lr=opt.LR,
momentum=opt.momentum,
weight_decay=opt.weightDecay)
else:
raise Exception
writer = SummaryWriter(
'.tensorboard/{}/{}'.format(opt.dataset, opt.expID)
)
# 准备训练数据
if opt.dataset == 'coco':
train_dataset =coco.Mscoco(train=True)
val_dataset = coco.Mscoco(train=False)
train_loader = torch.utils.data.DataLoader(
train_dataset, batch_size=opt.trainBatch, shuffle=True, num_workers=opt.nThreads, pin_memory=True
)
val_loader = torch.utils.data.DataLoader(
val_dataset, batch_size=opt.validBatch, shuffle=False, num_workers=opt.nThreads, pin_memory=True
)
m = torch.nn.DataParallel(m).cuda()
# 开始训练
for i in range(opt.nEpochs):
opt.epoch = i
print('############## String Epoch ##########'.format(opt.epoch))
loss, acc = train(train_loader, m, criterion, optimizer, writer)
print('Train-{idx:d} epoch | loss:{loss:.8f} | acc:{acc:.4f}'.format(
idx=opt.epoch,
loss=loss,
acc=acc
))
opt.acc = acc
opt.loss = loss
m_dev = m.module
if i % opt.snapshot == 0:
torch.save(
m_dev.state_dict(), '../exp/{}/{}/model_{}.pkl'.format(opt.dataset, opt.expID, opt.epoch)
)
torch.save(
opt, '../exp/{}/{}/option.pkl'.format(opt.dataset, opt.expID, opt.epoch)
)
torch.save(
optimizer, '../exp/{}/{}/optimizer.pkl'.format(opt.dataset, opt.expID)
)
loss, acc = valid(val_loader, m, criterion, optimizer, writer)
print(
'Valid-{idx:d} epoch | loss:{loss:.8f} | acc:{acc:.4f}'.format(
idx=i,
loss=loss,
acc=acc)
)
writer.close()
def main():
# Model Initialize
m = createModel().cuda()
if opt.loadModel:
print('Loading Model from {}'.format(opt.loadModel))
current_model_weight = m.state_dict()
weight_save = torch.load(opt.loadModel)
weight_save_changed = {}
for k in weight_save:
if 'conv_out.weight' in k or 'conv_out.bias' in k:
print(k, 'not used')
continue
weight_save_changed[k] = weight_save[k]
current_model_weight.update(weight_save_changed)
m.load_state_dict(current_model_weight)
if not os.path.exists("../exp/{}/{}".format(opt.dataset, opt.expID)):
try:
os.mkdir("../exp/{}/{}".format(opt.dataset, opt.expID))
except FileNotFoundError:
os.mkdir("../exp")
os.mkdir("../exp/{}".format(opt.dataset))
os.mkdir("../exp/{}/{}".format(opt.dataset, opt.expID))
else:
print('Create new model')
if not os.path.exists("../exp/{}/{}".format(opt.dataset, opt.expID)):
try:
os.mkdir("../exp/{}/{}".format(opt.dataset, opt.expID))
except FileNotFoundError:
os.mkdir("../exp")
os.mkdir("../exp/{}".format(opt.dataset))
os.mkdir("../exp/{}/{}".format(opt.dataset, opt.expID))
criterion = torch.nn.MSELoss().cuda()
if opt.optMethod == 'rmsprop':
optimizer = torch.optim.RMSprop(m.parameters(), lr=opt.LR)
elif opt.optMethod == 'adam':
optimizer = torch.optim.Adam(m.parameters(), lr=opt.LR)
else:
raise Exception
writer = SummaryWriter('.tensorboard/{}/{}'.format(opt.dataset, opt.expID))
# Prepare Dataset
if opt.dataset == 'coco':
train_dataset = coco.Mscoco(train=True,
img_folder=opt.img_folder_train,
annot_file=opt.annot_file_train,
nJoints=opt.nClasses)
val_dataset = coco.Mscoco(train=False,
img_folder=opt.img_folder_val,
annot_file=opt.annot_file_val,
nJoints=opt.nClasses)
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=opt.trainBatch,
shuffle=True,
num_workers=opt.nThreads,
pin_memory=True)
val_loader = torch.utils.data.DataLoader(val_dataset,
batch_size=opt.validBatch,
shuffle=False,
num_workers=opt.nThreads,
pin_memory=True)
# Model Transfer
m = DataParallel(m).cuda()
# Start Training
for i in range(opt.nEpochs + 1):
opt.epoch = i
print('############# Starting Epoch {} #############'.format(
opt.epoch))
loss, acc = train(train_loader, m, criterion, optimizer, writer)
print('Train-{idx:d} epoch | loss:{loss:.8f} | acc:{acc:.4f}'.format(
idx=opt.epoch, loss=loss, acc=acc))
opt.acc = acc
opt.loss = loss
m_dev = m.module
if (i % opt.snapshot == 0 and i != 0) or i == opt.nEpochs:
torch.save(
m_dev.state_dict(),
'../exp/{}/{}/model_{}.pkl'.format(opt.dataset, opt.expID,
opt.epoch))
torch.save(
opt, '../exp/{}/{}/option.pkl'.format(opt.dataset, opt.expID,
opt.epoch))
# torch.save(
# optimizer, '../exp/{}/{}/optimizer.pkl'.format(opt.dataset, opt.expID))
loss, acc = valid(val_loader, m, criterion, optimizer, writer)
print('Valid-{idx:d} epoch | loss:{loss:.8f} | acc:{acc:.4f}'.format(
idx=i, loss=loss, acc=acc))
# # Prediction Test
# if opt.map and acc > 0.8:
# with torch.no_grad():
# mAP, mAP5 = prediction_silence(m)
# logger.info('Prediction-{idx:d} epoch | mAP:{mAP:.3f} | mAP0.5:{mAP5:.3f}'.format(
# idx=i,
# mAP=mAP,
# mAP5=mAP5
# ))
'''
if opt.dataset != 'mpii':
with torch.no_grad():
mAP, mAP5 = prediction(m)
print('Prediction-{idx:d} epoch | mAP:{mAP:.3f} | mAP0.5:{mAP5:.3f}'.format(
idx=i,
mAP=mAP,
mAP5=mAP5
))
'''
writer.close()
