def train(params):
batch_size = 64
train_dataloader = torch.utils.data.DataLoader(ListDataset(['data/dataset/train.txt']), batch_size=batch_size,
shuffle=True, num_workers=0)
val_dataloader = torch.utils.data.DataLoader(ListDataset('data/dataset/validation.txt'), batch_size=batch_size,
shuffle=False,
num_workers=0)
model = make_model(len(char2token))
model.to(device=params['device'])
criterion = LabelSmoothing(size=len(char2token), padding_idx=0, smoothing=0.1)
criterion.to(device=params['device'])
model_opt = NoamOpt(model.tgt_embed[0].d_model, 1, 2000,
torch.optim.Adam(model.parameters(), lr=0, betas=(0.9, 0.98), eps=1e-9))
for epoch in range(10000):
params['epoch'] = epoch
model.train()
train_loss = run_epoch(train_dataloader, model,
SimpleLossCompute(model.generator, criterion, model_opt), params)
print("train_loss", train_loss)
model.eval()
val_loss = run_epoch(val_dataloader, model,
SimpleLossCompute(model.generator, criterion, None), params)
print("val_loss", val_loss)
torch.save(model.state_dict(), 'checkpoint/%08d_%f.pth' % (epoch, val_loss))
def train():
batch_size = 64
train_dataloader = torch.utils.data.DataLoader(ListDataset(
['your-train-lines']),
batch_size=batch_size,
shuffle=True,
num_workers=0)
val_dataloader = torch.utils.data.DataLoader(
ListDataset('your-test-lines'),
batch_size=batch_size,
shuffle=False,
num_workers=0)
model = make_model(len(char2token))
model.load_state_dict(torch.load('your-pretrain-model-path'))
model.cuda()
criterion = LabelSmoothing(size=len(char2token),
padding_idx=0,
smoothing=0.1)
criterion.cuda()
model_opt = NoamOpt(
model.tgt_embed[0].d_model, 1, 2000,
torch.optim.Adam(model.parameters(), lr=0, betas=(0.9, 0.98),
eps=1e-9))
for epoch in range(10000):
model.train()
run_epoch(train_dataloader, model,
SimpleLossCompute(model.generator, criterion, model_opt))
model.eval()
test_loss = run_epoch(
val_dataloader, model,
SimpleLossCompute(model.generator, criterion, None))
print("test_loss", test_loss)
torch.save(model.state_dict(),
'checkpoint/%08d_%f.pth' % (epoch, test_loss))
def insert_array_in_dataset(array, dt: ldt):
if dt is None:
return
if len(array) == 0 or array is None:
return
label = randomNumbers.generate_label()
dt.add_series(label, array)
def __init__(self, hparams):
super().__init__()
self.hparams = hparams
self.train_ds = ListDataset(hparams.train_ds, train=True)
self.valid_ds = ListDataset(hparams.valid_ds, train=False)
self.model = YOLOv4(n_classes=5, pretrained=hparams.pretrained, dropblock=hparams.Dropblock, sam=hparams.SAM, eca=hparams.ECA, ws=hparams.WS, iou_aware=hparams.iou_aware, coord=hparams.coord, hard_mish=hparams.hard_mish).cuda()
def __init__(self, hparams):
super().__init__()
self.hparams = hparams
self.train_ds = ListDataset(hparams.train_ds, train=True)
self.valid_ds = ListDataset(hparams.valid_ds, train=False)
self.model = YOLOv4(n_classes = 5, pretrained=True).cuda()
def __init__(self, path, img_size, batch_size, debug=False):
dataset = ListDataset(path,
img_size=img_size,
augment=False,
multiscale=False)
if debug:
dataset.img_files = dataset.img_files[:batch_size]
self.dataloader = torch.utils.data.DataLoader(
dataset,
batch_size=batch_size,
shuffle=False,
num_workers=0,
collate_fn=dataset.collate_fn)
def main(args):
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
args.device = device
print('device: {}'.format(device))
# create model
model = SPP_VGG16().to(device)
# load checkpoint
if args.model_weight:
if os.path.isfile(args.model_weight):
print("=> loading checkpoint '{}'".format(args.model_weight))
checkpoint = torch.load(args.model_weight, map_location=device)
model.load_state_dict(checkpoint['state_dict'])
print("=> loaded checkpoint '{}'".format(args.model_weight))
else:
print("=> no checkpoint found at '{}'".format(args.model_weight))
# Data loading code
test_path = cfg.TEST_PATH
test_label_path = cfg.TEST_LABEL_PATH
test_dataset = ListDataset(test_path, test_label_path, mode='test')
test_loader = torch.utils.data.DataLoader(
test_dataset,
batch_size=args.batch_size, shuffle=False,
num_workers=args.workers, pin_memory=True, collate_fn=test_dataset.collate_fn)
# Evaluate on test dataset
validate(test_loader, model, args)
def validate(val_list, model, criterion):
print('begin test')
test_loader = torch.utils.data.DataLoader(ListDataset(
val_list,
shuffle=False,
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]),
]),
train=False),
batch_size=args.batch_size)
model.eval()
mae = 0
for i, (img, target) in enumerate(test_loader):
img = img.cuda()
img = Variable(img)
output = model(img)
mae += abs(output.data.sum() -
target.sum().type(torch.FloatTensor).cuda())
mae = mae / len(test_loader)
print(' * MAE {mae:.3f} '.format(mae=mae))
return mae
def evaluate(model, transform, test_path, batch_sizes, thresh, im_size):
model.eval()
model.cuda()
the_classes = [
'crazing', 'inclusion', 'patches', 'pitted_surface', 'rolled-in_scale',
'scratches'
]
# test_path = '/home/ecust/gmy/pytorch-retinanet-master/pytorch-retinanet-master/data/NEU-DET/valid/labels'
testset = ListDataset(root=test_path,
train=False,
transform=transform,
input_size=im_size)
testloader = torch.utils.data.DataLoader(testset,
batch_size=batch_sizes,
shuffle=False,
num_workers=8,
collate_fn=testset.collate_fn)
encoder = DataEncoder()
evaluation(model, batch_sizes, testloader, the_classes, encoder, thresh,
im_size, im_size)
map = eval(the_classes)
return map
def train(train_list, model, criterion, optimizer, epoch):
losses = AverageMeter()
batch_time = AverageMeter()
data_time = AverageMeter()
train_loader = torch.utils.data.DataLoader(ListDataset(
train_list,
shuffle=True,
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]),
]),
train=True,
seen=model.seen,
batch_size=args.batch_size,
num_workers=args.workers),
batch_size=args.batch_size)
print('epoch %d, processed %d samples, lr %.10f' %
(epoch, epoch * len(train_loader.dataset), args.lr))
model.train()
end = time.time()
for i, (img, target) in enumerate(train_loader):
data_time.update(time.time() - end)
img = img.cuda()
img = Variable(img)
output = model(img)
target = target.type(torch.FloatTensor).unsqueeze(0).cuda()
target = Variable(target)
loss = criterion(output, target)
losses.update(loss.item(), img.size(0))
optimizer.zero_grad()
loss.backward()
optimizer.step()
batch_time.update(time.time() - end)
end = time.time()
if i % args.print_freq == 0:
print('Epoch: [{0}][{1}/{2}]\t'
'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
'Data {data_time.val:.3f} ({data_time.avg:.3f})\t'
'Loss {loss.val:.4f} ({loss.avg:.4f})\t'.format(
epoch,
i,
len(train_loader),
batch_time=batch_time,
data_time=data_time,
loss=losses))
def __init__(self, hparams):
super().__init__()
self.hparams = hparams
self.train_ds = ListDataset(hparams.train_ds, train=True, img_extensions=hparams.img_extensions)
self.valid_ds = ListDataset(hparams.valid_ds, train=False, img_extensions=hparams.img_extensions)
self.model = YOLOv4(n_classes=hparams.n_classes,
pretrained=hparams.pretrained,
dropblock=hparams.Dropblock,
sam=hparams.SAM,
eca=hparams.ECA,
ws=hparams.WS,
iou_aware=hparams.iou_aware,
coord=hparams.coord,
hard_mish=hparams.hard_mish,
asff=hparams.asff,
repulsion_loss=hparams.repulsion_loss,
acff=hparams.acff,
bcn=hparams.bcn,
mbn=hparams.mbn).cuda()
def train(args):
batch_size = 1
train_dataloader = torch.utils.data.DataLoader(ListDataset(
args.train_annot_path),
batch_size=batch_size,
shuffle=True,
num_workers=0)
val_dataloader = torch.utils.data.DataLoader(ListDataset(
args.test_annot_path),
batch_size=batch_size,
shuffle=False,
num_workers=0)
model = make_model(len(char2token))
if args.checkpoint:
model.load_state_dict(torch.load(args.checkpoint))
criterion = LabelSmoothing(size=len(char2token),
padding_idx=0,
smoothing=0.1)
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
model.cuda()
criterion.cuda()
model_opt = NoamOpt(
model.tgt_embed[0].d_model, 1, 2000,
torch.optim.Adam(model.parameters(), lr=0, betas=(0.9, 0.98),
eps=1e-9))
for epoch in range(100):
model.train()
run_epoch(train_dataloader,
model,
SimpleLossCompute(model.generator, criterion, model_opt),
count_iter=10000)
model.eval()
test_loss = run_epoch(
val_dataloader, model,
SimpleLossCompute(model.generator, criterion, None))
print("test_loss", test_loss)
torch.save(model.state_dict(),
'checkpoint/%08d_%f.pth' % (epoch, test_loss))
def tag(model: DeepEosModel,
text_file: Union[str, Path],
vocabulary_path: Union[str, Path],
batch_size=32,
window_size=5,
return_indices=False,
use_default_markers=True,
device=torch.device("cuda:0" if torch.cuda.is_available() else "cpu")):
with open(vocabulary_path, 'rb') as f:
vocabulary = pickle.load(f)
peos_list = []
with open(text_file, 'r', encoding='utf8') as f:
text = f.read()
peos_list.extend(
EosDataset.build_potential_eos_list(text, window_size,
use_default_markers))
text = np.asarray(list(text), dtype=str)
dataset = ListDataset(
EosDataset.build_data_set(peos_list, vocabulary, window_size))
ret_idx = []
model.eval()
dataloader = DataLoader(dataset, batch_size=batch_size)
with tqdm(dataloader, total=len(dataloader), desc="Tagging",
ascii=True) as tq:
for batch_no, (indices, x_tag) in enumerate(tq):
x_tag = x_tag.long().to(device)
prediction = model(x_tag) >= 0.5
zipped = zip(indices.squeeze().int().cpu().tolist(),
prediction.squeeze().cpu().tolist())
true_indices = [
idx + 1 for idx, _ in filter(lambda p: p[1], zipped)
]
if return_indices:
ret_idx.extend(true_indices)
else:
text[np.array(true_indices)] = '\n'
if return_indices:
return ret_idx
else:
return "".join(text)
net = FaceBox()
if use_gpu:
net.cuda()
print('load model...')
# net.load_state_dict(torch.load('weight/faceboxes.pt'))
criterion = MultiBoxLoss()
# optimizer = torch.optim.SGD(net.parameters(), lr=learning_rate, momentum=0.9, weight_decay=0.0003)
optimizer = torch.optim.Adam(net.parameters(),
lr=learning_rate,
weight_decay=1e-4)
train_dataset = ListDataset(root=file_root,
list_file='label/box_label.txt',
train=True,
transform=[transforms.ToTensor()])
train_loader = DataLoader(train_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=5)
print('the dataset has %d images' % (len(train_dataset)))
print('the batch_size is %d' % (batch_size))
num_iter = 0
vis = visdom.Visdom()
win = vis.line(Y=np.array([0]), X=np.array([0]))
net.train()
for epoch in range(num_epochs):
if epoch == 190 or epoch == 250:
def stepOnce(net, inputs, optimizer, netName="unknown", lossList=[], epoch=0):
dat, loss = net(*inputs)
loss.backward()
optimizer.step()
print("Epoch %d %s net training loss: %.5f" %
(epoch, netName, loss.item()))
lossList.append(loss.item())
return dat, loss
if __name__ == '__main__':
options = parse_args()
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
net = model.SuperResolNet(options)
net.to(device)
trainDataset = ListDataset(options.trainDataPath, options)
lossList = []
_trainCount = 0
if options.mode == "primaryT":
dataloader = DataLoader(trainDataset,
batch_size=options.batchSize,
shuffle=True,
num_workers=0,
pin_memory=True,
collate_fn=trainDataset.collate_fnSingleImg)
net.primaryNet.loadPretrainedParams("primaryParam.dict")
optimizer = optim.Adam(net.primaryNet.parameters(), lr=0.001, eps=1e-4)
for fEpoch in range(options.epochNum):
dataGen = iter(dataloader)
parser.add_argument('-c', '--classes_name', metavar='PATH',
default=cfg.CLASSES_NAME_PATH, type=str, help='path to classes name')
parser.add_argument('-w', '--model_weight', default=cfg.WEIGHTS_PATH, type=str, metavar='PATH',
help='path to latest checkpoint')
args = parser.parse_args()
print(args)
os.makedirs("./results", exist_ok=True)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print('device:{}'.format(device))
# load classes' name
classes_name = load_classes(args.classes_name)
# load data
dataset = ListDataset(args.data_path, args.bbox_path)
dataset_loader = torch.utils.data.DataLoader(
dataset, batch_size=1, num_workers=1,
pin_memory=True, collate_fn=dataset.collate_fn)
# create model
model = SPP_VGG16().to(device)
# load checkpoint
if args.model_weight:
if os.path.isfile(args.model_weight):
print("=> loading checkpoint '{}'".format(args.model_weight))
checkpoint = torch.load(args.model_weight, map_location=device)
model.load_state_dict(checkpoint['state_dict'])
else:
print("=> no checkpoint found at '{}'".format(args.model_weight))
batch_size = 1
data_encoder = DataEncoder()
print('Loading model..')
net = SSD300()
net.load_state_dict(torch.load('./checkpoint/ssd300_ckpt.pth')['net'])
net.to(device)
net.eval()
print('Preparing dataset..')
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225))
])
testset = ListDataset(root='/home/pzl/Data/tiny_vid',
list_file="/home/pzl/Data/tiny_vid/test_images.txt",
train=False,
transform=transform)
testloader = torch.utils.data.DataLoader(testset,
batch_size=batch_size,
shuffle=False,
num_workers=4)
acc = 0.
mean_iou = 0.
with torch.no_grad():
# batch_size = 1
for batch_idx, (images, loc_targets, conf_targets, real_boxes,
real_class) in enumerate(testloader):
# (1,3,300,300) cpu--> gpu
images = images.to(device)
# (1, 8732, 4) (1, 8732) cpu
def main(args):
best_acc1 = 0
os.makedirs('checkpoints', exist_ok=True)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
args.device = device
print('device: {}'.format(device))
# create model
model = SPP_VGG16().to(device)
model.apply(weights_init_normal)
# define optimizer
optimizer = torch.optim.Adam(model.parameters(), args.lr, weight_decay=args.weight_decay)
# optionally resume from a checkpoint
if args.resume:
if os.path.isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume, map_location=device)
args.start_epoch = checkpoint['epoch']
best_acc1 = checkpoint['best_acc1']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint '{}' (epoch {})"
.format(args.resume, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(args.resume))
# Data loading code
train_path = cfg.TRAIN_PATH
train_label_path = cfg.TRAIN_LABEL_PATH
val_path = cfg.VAL_PATH
val_label_path = cfg.VAL_LABEL_PATH
train_dataset = ListDataset(train_path, train_label_path)
print(len(train_dataset))
val_dataset = ListDataset(val_path, val_label_path)
train_loader = torch.utils.data.DataLoader(
train_dataset, batch_size=args.batch_size, shuffle=True,
num_workers=args.workers, pin_memory=True, collate_fn=train_dataset.collate_fn)
val_loader = torch.utils.data.DataLoader(
val_dataset, batch_size=args.batch_size, shuffle=False,
num_workers=args.workers, pin_memory=True, collate_fn=val_dataset.collate_fn)
logger = Logger('./logs')
for epoch in range(args.start_epoch, args.epochs):
# train for one epoch
adjust_learning_rate(optimizer, epoch, args)
train_loss, train_acc = train(train_loader, model, optimizer, epoch, args)
# evaluate on validation set
val_loss, val_acc = validate(val_loader, model, args)
# remember best acc@1 and save checkpoint
is_best = val_acc > best_acc1
best_acc1 = max(val_acc, best_acc1)
# log
info = {
'train_loss': float(train_loss),
'train_acc': float(train_acc),
'val_loss': float(val_loss),
'val_acc': float(val_acc)
}
for tag, value in info.items():
logger.scalar_summary(tag, value, epoch)
save_checkpoint({
'epoch': epoch + 1,
'state_dict': model.state_dict(),
'best_acc1': best_acc1,
'optimizer': optimizer.state_dict(),
}, is_best)
def __init__(self, args):
self.cfg = yaml.load(open('config.yaml', encoding='utf-8'))
self.cfg = EasyDict(self.cfg)
if args.model == 'vgg':
self.cfg = self.cfg.VGGConfig
if args.model == 'pse':
self.cfg = self.cfg.PSENetConfig
assert torch.cuda.is_available(), 'Error: CUDA not found!'
if not os.path.exists(args.save_folder):
os.mkdir(args.save_folder)
if not os.path.exists(args.logdir):
os.mkdir(args.logdir)
log_dir = args.logdir
if os.path.exists(log_dir):
shutil.rmtree(log_dir)
os.mkdir(log_dir)
# Data
print('==> Preparing data..')
self.trainset = ListDataset(
train=True,
transform=Augmentation_traininig,
input_size_min=args.input_size_min,
input_size_max=args.input_size_max,
multi_scale=args.multi_scale,
train_path='./test_data/image',
debug_num=4,
divisibility=self.cfg.DATA.SIZE_DIVISIBILITY,
step=self.cfg.ANCHOR_HEAD.STRIDE,
gen_heat_map=self.cfg.PIXEL_HEAD.PIXEL_PREDICTION)
self.trainloader = torch.utils.data.DataLoader(
self.trainset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
collate_fn=self.trainset.collate_fn,
)
print('==> Preparing data Done...')
# set model (focal_loss vs OHEM_CE loss)
self.ctpn_criterion = CTPNLoss()
self.seg_criterion = SegLoss()
self.num_classes = 1
# Training Detail option\
# if args.dataset in ["SynthText"] else (2000, 4000, 6000, 8000, 10000)
self.stepvalues = (10000, 20000, 30000, 40000, 50000)
self.best_loss = float('inf') # best test loss
self.start_epoch = 0 # start from epoch 0 or last epoch
self.iteration = 0
self.cur_lr = args.lr
self.mean = (0.485, 0.456, 0.406)
self.var = (0.229, 0.224, 0.225)
self.step_index = 0
self.pEval = None
# Model
# 注意： resnet中， downsample 部分会用 stride = (2,2)的 1x1 conv ，会使得 25 x 25 -> 13 x 13 ， 和 anchor 处理机制不符
# 所以resnet 中，一定要对数据做 SIZE_DIVISIBILITY
self.net = CTPN(self.cfg) # plus one means the bg
if args.resume:
# 这里要做模型加载预训练模型，在finetune是使用
print('==> Resuming from checkpoint..', args.resume)
checkpoint = torch.load(args.resume)
self.net.load_state_dict(checkpoint['net'])
print("multi_scale : ", args.multi_scale)
print("input_size : ", args.input_size_min)
print("stepvalues : ", self.stepvalues)
print("start_epoch : ", self.start_epoch)
print("iteration : ", self.iteration)
print("cur_lr : ", self.cur_lr)
print("step_index : ", self.step_index)
print("num_gpus : ", torch.cuda.device_count())
self.net = torch.nn.DataParallel(self.net,
device_ids=list(
range(torch.cuda.device_count())))
self.net.cuda()
# Training
self.net.train()
# net.module.freeze_bn() # you must freeze batchnorm
# optimizer = optim.SGD(net.parameters(), lr=cur_lr, momentum=0.9, weight_decay=1e-4)
self.optimizer = optim.Adam(self.net.parameters(), lr=self.cur_lr)
# encoder = DataEncoder()
self.encoder = self.trainset.encoder
pickle.dump(self.encoder,
open(os.path.join(args.save_folder, 'encoder.pkl'), 'wb'))
# tensorboard visualize
self.writer = SummaryWriter(log_dir=args.logdir)
def generate_empty_dataset():
return ldt.ListDataset()
# print('load model...')
# net.load_state_dict(torch.load('weight/faceboxes.pt'))
criterion = MultiBoxLoss()
# optimizer = torch.optim.SGD(net.parameters(), lr=learning_rate, momentum=0.9, weight_decay=0.0005)
optimizer = torch.optim.Adam(net.parameters(),
lr=learning_rate,
weight_decay=1e-5)
# optimizer = torch.optim.RMSprop(net.parameters(), lr=learning_rate, alpha=0.9)
train_dataset = ListDataset(root=file_root,
data_type=data_type,
list_file=label_path,
lmdb_path=lmdb_path,
train=True,
transform=[transforms.ToTensor()],
multi_scale=multi_scale)
if not multi_scale:
train_loader = DataLoader(train_dataset,
batch_size=batch_size,
shuffle=True,
pin_memory=True,
num_workers=12)
else:
train_loader = DataLoader(train_dataset,
batch_sampler=BatchSampler(
RandomSampler(train_dataset),
batch_size,
True,
def train():
use_gpu = torch.cuda.is_available()
file_root = os.path.dirname(os.path.abspath(__file__))
learning_rate = 0.001
num_epochs = 300
batch_size = 32
net = FaceBox()
if use_gpu:
net.cuda()
print('load model...')
net.load_state_dict(torch.load('weight/faceboxes_add_norm.pt'))
criterion = MultiBoxLoss()
#optimizer = torch.optim.SGD(net.parameters(), lr=learning_rate, momentum=0.9, weight_decay=0.0005)
optimizer = torch.optim.Adam(net.parameters(), lr=learning_rate, weight_decay=1e-4)
scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[198, 248], gamma=0.1)
train_dataset = ListDataset(root=file_root,
list_file='data/train_rewrite.txt',
train=True,
transform = [transforms.ToTensor()])
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=4)
val_dataset = ListDataset(root=file_root,
list_file='data/val_rewrite.txt',
train=False,
transform = [transforms.ToTensor()])
val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=4)
print('the dataset has %d images' % (len(train_dataset)))
print('the batch_size is %d' % (batch_size))
num_iter = 0
vis = visdom.Visdom()
win = vis.line(Y=np.array([0]), X = np.array([0]))
net.train()
for epoch in range(num_epochs):
scheduler.step()
print('\n\nStarting epoch %d / %d' % (epoch + 1, num_epochs))
print('Learning Rate for this epoch: {}'.format(learning_rate))
total_loss = 0.
net.train()
for i,(images,loc_targets,conf_targets) in enumerate(train_loader):
if use_gpu:
images = images.cuda()
loc_targets = loc_targets.cuda()
conf_targets = conf_targets.cuda()
loc_preds, conf_preds = net(images)
loss = criterion(loc_preds,loc_targets,conf_preds,conf_targets)
total_loss += loss.item()
optimizer.zero_grad()
loss.backward()
optimizer.step()
if (i+1) % 10 == 0:
print ('Epoch [{}/{}], Iter [{}/{}] Loss: {:.4f}, average_loss: {:.4f}'.format(
epoch+1, num_epochs, i+1, len(train_loader), loss.item(), total_loss / (i+1)))
#train_loss = total_loss /(len(train_dataset) / batch_size)
vis.line(Y=np.array([total_loss / (i+1)]), X=np.array([num_iter]),
win=win,
name='train',
update='append')
num_iter += 1
# val_loss = 0.0
# net.eval()
# for idx, (images, loc_targets,conf_targets) in enumerate(val_loader):
# with torch.no_grad():
# if use_gpu:
# images = images.cuda()
# loc_targets = loc_targets.cuda()
# conf_targets = conf_targets.cuda()
#
# loc_preds, conf_preds = net(images)
# loss = criterion(loc_preds, loc_targets, conf_preds, conf_targets)
# val_loss += loss.item()
# val_loss /= len(val_dataset)/batch_size
# vis.line(Y=np.array([val_loss]), X=np.array([epoch]),
# win=win,
# name='val',
# update='append')
# print('loss of val is {}'.format(val_loss))
if not os.path.exists('weight/'):
os.mkdir('weight')
print('saving model ...')
torch.save(net.state_dict(),'weight/faceboxes_add_norm.pt')
def main(argv):
parser = argparse.ArgumentParser()
parser.add_argument("--epochs",
type=int,
default=10000,
help="number of epochs")
parser.add_argument("--batch_size",
type=int,
default=20,
help="size of each image batch")
parser.add_argument("--data_config",
type=str,
default="config/adc.data",
help="path to data config file")
# parser.add_argument("--pretrained_weights", type=str, default="config/yolov3_ckpt_5.pth") # models/model1/yolov3_ckpt_73.pth
parser.add_argument("--pretrained_weights",
type=str) # models/model1/yolov3_ckpt_73.pth
parser.add_argument(
"--n_cpu",
type=int,
default=0,
help="number of cpu threads to use during batch generation")
parser.add_argument("--img_size",
type=int,
default=[768, 1024],
help="size of each image dimension")
parser.add_argument("--evaluation_interval",
type=int,
default=1,
help="interval evaluations on validation set")
parser.add_argument("--multiscale",
default='False',
choices=['True', 'False'])
parser.add_argument("--augment",
default='False',
choices=['True', 'False'])
parser.add_argument("--save_path",
type=str,
default='models/weights_1350_0102',
help="save model path")
parser.add_argument("--debug",
type=str,
default='False',
choices=['True', 'False'],
help="debug")
parser.add_argument("--lr", type=float, default=0.01, help="learning rate")
args = parser.parse_args(argv)
args.debug = True if args.debug == 'True' else False
args.multiscale = True if args.multiscale == 'True' else False
args.augment = True if args.augment == 'True' else False
print_args(args)
print(
datetime.datetime.strftime(datetime.datetime.now(),
'%Y-%m-%d %H:%M:%S'))
if args.debug:
print('debug...')
import shutil
# if os.path.exists(args.save_path):
# shutil.rmtree(args.save_path)
args.evaluation_interval = 1
# debug模式下先删除save_path,并每间隔一轮验证一次
# assert not os.path.exists(args.save_path)
# os.makedirs(args.save_path)
# adc.dat下有train和valid两个dat还有anchor.txt的路径
data_config = parse_data_config(args.data_config)
train_path = data_config["train"]
valid_path = data_config["valid"]
if args.debug:
valid_path = train_path
anchors = get_anchors(data_config['anchors']).to('cuda')
model = ResNet(anchors).to('cuda')
if args.pretrained_weights:
print('pretrained weights: ', args.pretrained_weights)
model.load_pretrained_weights(args.pretrained_weights)
dataset = ListDataset(train_path,
img_size=args.img_size,
augment=args.augment,
multiscale=args.multiscale)
eval = evaluate(path=valid_path,
img_size=args.img_size,
batch_size=args.batch_size,
debug=args.debug)
if args.debug:
dataset.img_files = dataset.img_files[:10 * args.batch_size]
dataloader = torch.utils.data.DataLoader(
dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.n_cpu,
collate_fn=dataset.collate_fn,
)
print('Number train sample: ', len(dataset))
optimizer = torch.optim.Adam(model.parameters(),
lr=args.lr,
weight_decay=5e-5)
# 这里优化器和学习率是不是要调节？
print('\n### train ...')
for epoch in range(args.epochs):
model.train()
lr = max(1e-10, args.lr * (0.95**epoch))
for param_group in optimizer.param_groups:
param_group['lr'] = lr
for batch_i, (imgs, targets, _) in enumerate(dataloader):
imgs = Variable(imgs.to('cuda'))
# 训练集有经过augment_sequential，而验证集没有
# targets=([[0.0000, 0.7328, 0.2808, 0.0934, 0.0808],
# [1.0000, 0.5255, 0.5466, 0.0596, 0.1587],
# [1.0000, 0.5585, 0.8077, 0.0553, 0.2250],
# [3.0000, 0.4519, 0.4351, 0.1365, 0.2048]], device='cuda:0')
targets = Variable(targets.to('cuda'), requires_grad=False)
yolo_map, _ = model(imgs)
# yolo_map.shape : [4,] 其中每个yolo_map的格式如下： batch,featuremap_h,featuremap_w,anchor_num,(x,y,w,h,conf)
loss, metrics = model.loss(yolo_map, targets)
loss.backward()
optimizer.step()
optimizer.zero_grad()
if (batch_i + 1) % 100 == 0 or (batch_i + 1) == len(dataloader):
time_str = datetime.datetime.strftime(datetime.datetime.now(),
'%Y-%m-%d %H:%M:%S')
lr = optimizer.param_groups[0]['lr']
loss = metrics["loss"]
xy = metrics["xy"]
wh = metrics["wh"]
conf = metrics["conf"]
loss_str = 'loss: {:<8.2f}'.format(loss)
loss_str += 'xy: {:<8.2f}'.format(xy)
loss_str += 'wh: {:<8.2f}'.format(wh)
loss_str += 'conf: {:<8.2f}'.format(conf)
epoch_str = 'Epoch: {:4}({:4}/{:4})'.format(
epoch, batch_i + 1, len(dataloader))
print('[{}]{} {} lr:{}'.format(time_str, epoch_str, loss_str,
lr))
print()
if epoch % args.evaluation_interval == 0:
print("\n---- Evaluating Model ----")
save_model_epoch = 'yolov3_ckpt_{}.pth'.format(epoch)
model.save_weights(os.path.join(args.save_path, save_model_epoch))
print(save_model_epoch)
example_save_path = args.save_path
for conf in [0.1, 0.3, 0.5, 0.7]:
metrics = eval(model,
iou_thres=0.5,
conf_thres=conf,
nms_thres=0.5,
save_path=example_save_path)
example_save_path = None
print(
'image_acc: {}\t{}\tbbox_acc: {}\tbbox_recall: {}'.format(
*metrics[1:]))
names = ['image', 'ture', 'det', 'box_acc', 'image_acc']
print('{:<10}{:<10}{:<10}{:<10}{:<10}'.format(*names))
print('{:<10}{:<10}{:<10}{:<10}{:<10}'.format(*metrics[0][0]))
print('{:<10}{:<10}{:<10}{:<10}{:<10}'.format(*metrics[0][1]))
print()
def __init__(self, train_path, test_path, model_file, model, img_size=1024, batch_size=16, lr=1e-3,
re_train=False, best_loss=2, use_gpu=False, nms_threshold=0.5):
self.train_path = train_path
self.test_path = test_path
self.model_file = model_file
self.img_size = img_size
self.batch_size = batch_size
self.re_train = re_train # 不加载训练模型，重新进行训练
self.best_loss = best_loss # 最好的损失值，小于这个值，才会保存模型
self.use_gpu = False
self.nms_threshold = nms_threshold
if use_gpu is True:
print("gpu available: %s" % str(torch.cuda.is_available()))
if torch.cuda.is_available():
self.use_gpu = True
else:
self.use_gpu = False
# 模型
self.model = model
if self.use_gpu:
print('[use gpu] ...')
self.model = self.model.cuda()
# 加载模型
if os.path.exists(self.model_file) and not self.re_train:
self.load(self.model_file)
# RandomHorizontalFlip
self.transform_train = T.Compose([
# T.Resize((self.img_size, self.img_size)),
T.ToTensor(),
# T.Normalize(mean=[.5, .5, .5], std=[.5, .5, .5]),
])
self.transform_test = T.Compose([
# T.Resize((self.img_size, self.img_size)),
T.ToTensor(),
# T.Normalize(mean=[.5, .5, .5], std=[.5, .5, .5])
])
# Dataset
train_label = os.path.join(self.train_path, 'faceboxes_label.txt')
test_label = os.path.join(self.test_path, 'faceboxes_label.txt')
train_dataset = ListDataset(root=self.train_path, list_file=train_label, train=True, transform=self.transform_train)
self.train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=1)
test_dataset = ListDataset(root=self.test_path, list_file=test_label, train=False, transform=self.transform_test)
self.test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=1)
print('train_loader len: %d' % len(self.train_loader.dataset))
print(' test_loader len: %d' % len(self.test_loader.dataset))
self.criterion = MultiBoxLoss()
self.lr = lr
# self.optimizer = optim.SGD(self.model.parameters(), lr=self.lr, momentum=0.5)
self.optimizer = optim.Adam(self.model.parameters(), lr=self.lr, weight_decay=1e-4)
# optimizer = torch.optim.Adam(net.parameters(), lr=learning_rate, weight_decay=1e-4)
pass
4: "deer",
5: "dog",
6: "frog",
7: "horse",
8: "ship",
9: "truck"
}
CIFAR_MEAN = [125.3 / 255, 123.0 / 255, 113.9 / 255]
if __name__ == '__main__':
options = parse_args()
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
net = model.AutoEncClassifier(options)
net.to(device)
trainDataset = ListDataset(options.imageSize,
options.trainImL,
options.trainLabL,
setRandom=True)
testDataset = ListDataset(options.imageSize, options.testImL,
options.testLabL)
trainDataloader = DataLoader(trainDataset,
batch_size=options.batchSize,
shuffle=True,
num_workers=2,
pin_memory=True,
collate_fn=trainDataset.collate_fn)
testDataloader = DataLoader(testDataset,
batch_size=200,
shuffle=False,
num_workers=2,
from tqdm import tqdm
from config import cfg
from dataset import ListDataset
from model import FasterRCNN
from torch.utils.data import DataLoader
from utils import array_tool as at
from utils.eval_tool import Eval
from terminaltables import AsciiTable
import visdom
import numpy as np
import torch
if __name__ == '__main__':
# 准备训练与验证数据
trainset = ListDataset(cfg.train_dir)
dataloader = DataLoader(trainset, batch_size=1, shuffle=True,num_workers=cfg.num_workers)
testset = ListDataset(cfg.val_dir, is_train=False)
test_dataloader = DataLoader(testset,batch_size=1,num_workers=cfg.test_num_workers,shuffle=False,pin_memory=True)
# 加载模型与权重
model = FasterRCNN().cuda()
if cfg.load_path:
model.load(cfg.load_path)
print('已加载训练模型')
# 创建visdom可视化端口
vis = visdom.Visdom(env='Faster R-CNN')
best_map = 0
for epoch in range(1,cfg.epoch):
for img, target_box, target_label, scale in tqdm(dataloader):
scale = at.scalar(scale)
img, target_box, target_label = img.cuda().float(), target_box.cuda(), target_label.cuda()
type=int,
default=384,
help='Output image size. Output an Hr image')
parser.add_argument('--inImgPatchSize',
type=int,
default=384,
help='Input image patch size')
parser.add_argument('--batchSize', type=int, default=2, help='Batch size')
args = parser.parse_args()
return args
if __name__ == '__main__':
options = parse_args()
testDataset = ListDataset(options.testDataPath, options)
with open(options.normCsvFile) as fil:
strr = fil.read()
strL = strr.split("\n")
dic = {f.split(" ")[0]: float(f.split(" ")[1]) for f in strL}
allSceneList = []
allNormList = []
for fInd in range(len(testDataset)):
imO = testDataset[fInd]
print(imO.sceneDir.split("\\")[-1])
norm = dic[imO.sceneDir.split("\\")[-1]]
_a, _b = imO()
Hr = torch.cat([_a, _b], dim=0).unsqueeze(0)
if torch.mean(_b) <= 0.8: continue
from model.Retinanet import RetinaNet
from dataset import ListDataset
import torch
import visdom
from tqdm import tqdm
from config import cfg
from utils.eval_tools import Eval
from utils.loss import multiboxloss
from terminaltables import AsciiTable
from torch.utils.data import DataLoader
train_dataset = ListDataset(cfg.train_dir, is_train=True)
test_dataset = ListDataset(cfg.val_dir, is_train=False)
if __name__ == '__main__':
data_loader = DataLoader(train_dataset,
batch_size=cfg.batch_size,
shuffle=True,
num_workers=cfg.num_workers)
model = RetinaNet(cfg.res_name).cuda()
if cfg.use_adam:
optimizer = torch.optim.Adam(model.parameters(),
lr=cfg.lr,
weight_decay=cfg.weight_decay)
else:
optimizer = torch.optim.SGD(model.parameters(),
lr=cfg.lr,
momentum=0.9,
weight_decay=cfg.weight_decay)
if cfg.pretrain:
model.load_state_dict(torch.load(cfg.load_path))