help='print frequency (default: 10)')
args = parser.parse_args()
# ##############################################################################################################
# ################################### Setup for some configurations ###########################################
# ##############################################################################################################
# 如果我们每次训练的输入数据的size不变，那么开启这个就会加快我们的训练速度
torch.backends.cudnn.benchmark = True
use_cuda = torch.cuda.is_available()
checkpoint_path = args.checkpoint_path
# > TOCHECK: training configs
opt = TrainingOpt()
config = GetConfig(opt.config_name)
soureconfig = COCOSourceConfig(opt.hdf5_train_data) # > 512.h5
train_data = MyDataset(config, soureconfig, shuffle=False,
augment=True) # shuffle in data loader
soureconfig_val = COCOSourceConfig(opt.hdf5_val_data)
val_data = MyDataset(config, soureconfig_val, shuffle=False,
augment=False) # shuffle in data loader
best_loss = float('inf')
start_epoch = 0 # 从0开始或者从上一个epoch开始
args.distributed = False
if 'WORLD_SIZE' in os.environ:
args.distributed = int(os.environ['WORLD_SIZE']) > 1
args.gpu = 0
args.world_size = 1
def train(**kwargs):
opt._parse(kwargs)
# device_num = 6
data_root = "/home/lsm/TrainSet/"
train_file = "train.txt"
test_file = "test.txt"
trainset = MyDataset(data_root, train_file, opt)
testset = TestDataset(data_root, test_file, opt)
print('load data')
dataloader = data_.DataLoader(trainset, \
batch_size=1, \
shuffle=True, \
# pin_memory=True,
num_workers=opt.num_workers)
test_dataloader = data_.DataLoader(testset,
batch_size=1,
num_workers=opt.test_num_workers,
shuffle=False, \
pin_memory=True
)
faster_rcnn = FasterRCNNVGG16()
print('model construct completed')
trainer = FasterRCNNTrainer(faster_rcnn).cuda()
if opt.load_path:
trainer.load(opt.load_path)
print('load pretrained model from %s' % opt.load_path)
# trainer.vis.text(dataset.db.label_names, win='labels')
best_map = 0
lr_ = opt.lr
f = open('log.txt', 'w')
for epoch in range(opt.epoch):
trainer.reset_meters()
print("epoch " + str(epoch) + " ...")
for ii, (img, bbox_, label_, scale) in tqdm(enumerate(dataloader)):
# break
# for ii, (img, bbox_, label_, scale) in enumerate(dataloader):
# print(ii)
scale = at.scalar(scale)
img, bbox, label = img.cuda().float(), bbox_.cuda(), label_.cuda()
trainer.train_step(img, bbox, label, scale)
# if (ii + 1) % opt.plot_every == 0:
# if os.path.exists(opt.debug_file):
# ipdb.set_trace()
#
# # plot loss
# # trainer.vis.plot_many(trainer.get_meter_data())
#
# # plot groud truth bboxes
# ori_img_ = inverse_normalize(at.tonumpy(img[0]))
# # gt_img = visdom_bbox(ori_img_,
# # at.tonumpy(bbox_[0]),
# # at.tonumpy(label_[0]))
# # trainer.vis.img('gt_img', gt_img)
#
# # plot predicti bboxes
# _bboxes, _labels, _scores = trainer.faster_rcnn.predict([ori_img_], visualize=True)
eval_result = eval(test_dataloader, faster_rcnn, test_num=opt.test_num)
lr_ = trainer.faster_rcnn.optimizer.param_groups[0]['lr']
log_info = 'lr:{}, map:{},loss:{}'.format(
str(lr_), str(eval_result['map']), str(trainer.get_meter_data()))
print(log_info)
# trainer.vis.log(log_info)
# print(str(lr_)+": loss = "+str(trainer.get_meter_data()))
f.write(log_info)
if eval_result['map'] > best_map:
best_map = eval_result['map']
best_path = trainer.save(best_map=best_map)
if epoch == 9:
trainer.load(best_path)
trainer.faster_rcnn.scale_lr(opt.lr_decay)
lr_ = lr_ * opt.lr_decay
if epoch == 19:
trainer.load(best_path)
trainer.faster_rcnn.scale_lr(opt.lr_decay)
lr_ = lr_ * opt.lr_decay
if epoch == 50:
break
f.close()
import torch
import os
from torchvision import transforms, models
from torch.utils.data import DataLoader
######pic_size = 480
#######pic_resize =224
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
if os.path.exists('./picture_total/testset/test.txt'):
print('test.txt has been existed')
transforms = transforms.Compose(
[transforms.Resize(224), transforms.ToTensor()])
test_data = MyDataset(txt='./picture_total/testset/test.txt',
transform=transforms)
test_loader = DataLoader(dataset=test_data, batch_size=10, num_workers=8)
else:
print('you need to prepare your test.txt first!')
# 只加载训练好的参数
def test_PalmLocNet(test_x, test_y):
vgg = models.vgg16_bn(pretrained=False)
PLNet = vggPalmLocNet(vgg)
PLNet.eval()
if torch.cuda.is_available():
PLNet.load_state_dict({
k.replace('module.', ''): v
for k, v in torch.load(
'./checkpoints/train_params_best.pth').items()
help="folder to store model")
parser.add_argument('-p','--PICTUREFOLDER',type= str, default='./picture_total/',
help="folder to store trained picture")
args = parser.parse_args()
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
if os.path.exists(args.PICTUREFOLDER+'trainset/'+'train.txt') and os.path.exists(args.PICTUREFOLDER+'testset/'+'test.txt') :
print('train.txt and test.txt have been existed')
transforms = transforms.Compose([
transforms.Resize(224),
transforms.ToTensor()
])
train_data = MyDataset(txt=args.PICTUREFOLDER + 'trainset/' + 'train.txt', transform=transforms)
test_data = MyDataset(txt=args.PICTUREFOLDER + 'testset/' + 'test.txt', transform=transforms)
train_loader = DataLoader(dataset=train_data, batch_size=args.BATCH_SIZE, shuffle=True, num_workers=8)
test_loader = DataLoader(dataset=test_data, batch_size=10, num_workers=8)
else:
print('you need to prepare your train.txt and test.txt first!')
#测试数据
for k, (tx, ty) in enumerate(test_loader):
test_x = tx.to(device)
test_y = ty.to(device)/480
# 训练以及保存模型数据
def train_PalmLocNet(train_loader, test_x, test_y):
if os.path.exists(args.MODELFOLDER + 'train_params_best.pth'):
print('reload the last best model parameters')