def main():
print('ASL Example Inference code on a single image')
# parsing args
args = parse_args(parser)
# setup model
print('creating and loading the model...')
state = torch.load(args.model_path, map_location='cpu')
args.num_classes = state['num_classes']
print('num_classes: ', args.num_classes)
classes_list = np.array(list(state['idx_to_class'].values()))
print('classes_list: ', classes_list)
model = create_model(args).cuda()
model.load_state_dict(state['model'], strict=True)
model.eval()
print('sigmoid threshold: ', args.th)
# doing inference
print('loading image and doing inference...')
im = Image.open(args.pic_path)
im_resize = im.resize((args.input_size, args.input_size))
np_img = np.array(im_resize, dtype=np.uint8)
tensor_img = torch.from_numpy(np_img).permute(
2, 0, 1).float() / 255.0 # HWC to CHW
tensor_batch = torch.unsqueeze(tensor_img, 0).cuda()
output = torch.squeeze(torch.sigmoid(model(tensor_batch)))
np_output = output.cpu().detach().numpy()
detected_classes = classes_list[np_output > args.th]
print('done\n')
# example loss calculation
output = model(tensor_batch)
loss_func1 = AsymmetricLoss()
loss_func2 = AsymmetricLossOptimized()
target = output.clone()
target[output < 0] = 0 # mockup target
target[output >= 0] = 1
loss1 = loss_func1(output, target)
loss2 = loss_func2(output, target)
assert abs((loss1.item() - loss2.item())) < 1e-6
# displaying image
print('showing image on screen...')
fig = plt.figure()
plt.imshow(im)
plt.axis('off')
plt.axis('tight')
# plt.rcParams["axes.titlesize"] = 10
plt.title("detected classes: {}".format(detected_classes))
plt.show()
print('done\n')
def main():
args = parser.parse_args()
args.batch_size = args.batch_size
# setup model
print('creating model...')
#state = torch.load(args.model_path, map_location='cpu')
#args.num_classes = state['num_classes']
args.do_bottleneck_head = True
model = create_model(args).cuda()
ema = EMA(model, 0.999)
ema.register()
#model.load_state_dict(state['model'], strict=True)
#model.train()
classes_list = np.array(list(idx_to_class.values()))
print('done\n')
# Data loading code
normalize = transforms.Normalize(mean=[0, 0, 0], std=[1, 1, 1])
instances_path_val = os.path.join(args.data,
'annotations/instances_val2017.json')
#instances_path_train = os.path.join(args.data, 'annotations/instances_val2017.json')#temprarily use val as train
instances_path_train = os.path.join(
args.data, 'annotations/instances_train2017.json')
data_path_val = os.path.join(args.data, 'val2017')
#data_path_train = os.path.join(args.data, 'val2017')#temporarily use val as train
data_path_train = os.path.join(args.data, 'train2017')
val_dataset = CocoDetection(
data_path_val, instances_path_val,
transforms.Compose([
transforms.Resize((args.image_size, args.image_size)),
transforms.ToTensor(),
normalize,
]))
train_dataset = CocoDetection(
data_path_train, instances_path_train,
transforms.Compose([
transforms.Resize((args.image_size, args.image_size)),
transforms.ToTensor(),
normalize,
]))
print("len(val_dataset)): ", len(val_dataset))
print("len(train_dataset)): ", len(train_dataset))
val_loader = torch.utils.data.DataLoader(val_dataset,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=False)
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=False)
criterion = AsymmetricLoss()
params = model.parameters()
optimizer = torch.optim.Adam(params, lr=0.0002,
weight_decay=0.0001) #尝试新的optimizer
total_step = len(train_loader)
scheduler = lr_scheduler.OneCycleLR(optimizer,
max_lr=0.0002,
total_steps=total_step,
epochs=25)
#total_step = len(train_loader)
highest_mAP = 0
trainInfoList = []
Sig = torch.nn.Sigmoid()
#f=open('info_train.txt', 'a')
for epoch in range(5):
for i, (inputData, target) in enumerate(train_loader):
f = open('info_train.txt', 'a')
#model.train()
inputData = inputData.cuda()
target = target.cuda()
target = target.max(dim=1)[0]
#Sig = torch.nn.Sigmoid()
output = Sig(model(inputData))
#output[output<args.thre] = 0
#output[output>=args.thre]=1
#print(output.shape) #(batchsize, channel, imhsize, imgsize)
#print(inputData.shape) #(batchsize, numclasses)
#print(output[0])
#print(target[0])
loss = criterion(output, target)
model.zero_grad()
loss.backward()
optimizer.step()
ema.update()
#store information
if i % 10 == 0:
trainInfoList.append([epoch, i, loss.item()])
print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'.format(
epoch, 5, i, total_step, loss.item()))
f.write('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}\n'.format(
epoch, 5, i, total_step, loss.item()))
if (i + 1) % 400 == 0:
#储存相应迭代模型
torch.save(
model.state_dict(),
os.path.join('models/',
'model-{}-{}.ckpt'.format(epoch + 1, i + 1)))
#modelName = 'models/' + 'decoder-{}-{}.ckpt'.format(epoch+1, i+1)
mAP_score = validate_multi(val_loader, model, args, ema)
#model.train()
if mAP_score > highest_mAP:
highest_mAP = mAP_score
print('current highest_mAP = ', highest_mAP)
f.write('current highest_mAP = {}\n'.format(highest_mAP))
torch.save(model.state_dict(),
os.path.join('models/', 'model-highest.ckpt'))
f.close()
scheduler.step() #修改学习率
def train_multi_label_coco(model, train_loader, val_loader, lr):
ema = ModelEma(model, 0.9997) # 0.9997^641=0.82
# set optimizer
Epochs = 80
Stop_epoch = 40
weight_decay = 1e-4
criterion = AsymmetricLoss(gamma_neg=4, gamma_pos=0, clip=0.05, disable_torch_grad_focal_loss=True)
parameters = add_weight_decay(model, weight_decay)
optimizer = torch.optim.Adam(params=parameters, lr=lr, weight_decay=0) # true wd, filter_bias_and_bn
steps_per_epoch = len(train_loader)
scheduler = lr_scheduler.OneCycleLR(optimizer, max_lr=lr, steps_per_epoch=steps_per_epoch, epochs=Epochs,
pct_start=0.2)
highest_mAP = 0
trainInfoList = []
scaler = GradScaler()
for epoch in range(Epochs):
if epoch > Stop_epoch:
break
for i, (inputData, target) in enumerate(train_loader):
inputData = inputData.cuda()
target = target.cuda() # (batch,3,num_classes)
target = target.max(dim=1)[0]
with autocast(): # mixed precision
output = model(inputData).float() # sigmoid will be done in loss !
loss = criterion(output, target)
model.zero_grad()
scaler.scale(loss).backward()
# loss.backward()
scaler.step(optimizer)
scaler.update()
# optimizer.step()
scheduler.step()
ema.update(model)
# store information
if i % 100 == 0:
trainInfoList.append([epoch, i, loss.item()])
print('Epoch [{}/{}], Step [{}/{}], LR {:.1e}, Loss: {:.1f}'
.format(epoch, Epochs, str(i).zfill(3), str(steps_per_epoch).zfill(3),
scheduler.get_last_lr()[0], \
loss.item()))
try:
torch.save(model.state_dict(), os.path.join(
'models/', 'model-{}-{}.ckpt'.format(epoch + 1, i + 1)))
except:
pass
model.eval()
mAP_score = validate_multi(val_loader, model, ema)
model.train()
if mAP_score > highest_mAP:
highest_mAP = mAP_score
try:
torch.save(model.state_dict(), os.path.join(
'models/', 'model-highest.ckpt'))
except:
pass
print('current_mAP = {:.2f}, highest_mAP = {:.2f}\n'.format(mAP_score, highest_mAP))
def train_multi_label_coco(model, ema, train_loader, val_loader, lr, args, logger):
# set optimizer
Epochs = 40
weight_decay = 1e-4
criterion = AsymmetricLoss(gamma_neg=4, gamma_pos=0, clip=0.05,
disable_torch_grad_focal_loss=args.dtgfl)
parameters = add_weight_decay(model, weight_decay)
optimizer = torch.optim.Adam(params=parameters, lr=lr, weight_decay=0) # true wd, filter_bias_and_bn
steps_per_epoch = len(train_loader)
scheduler = lr_scheduler.OneCycleLR(optimizer, max_lr=lr, steps_per_epoch=steps_per_epoch, epochs=Epochs, pct_start=0.2)
highest_mAP = 0
trainInfoList = []
scaler = GradScaler()
for epoch in range(Epochs):
for i, (inputData, target) in enumerate(train_loader):
# break
inputData = inputData.cuda()
target = target.cuda() # (batch,3,num_classes)
target = target.max(dim=1)[0]
with autocast(): # mixed precision
output = model(inputData).float() # sigmoid will be done in loss !
loss = criterion(output, target)
model.zero_grad()
scaler.scale(loss).backward()
# loss.backward()
scaler.step(optimizer)
scaler.update()
# optimizer.step()
scheduler.step()
ema.update(model)
# store information
if i % 100 == 0:
trainInfoList.append([epoch, i, loss.item()])
logger.info('Epoch [{}/{}], Step [{}/{}], LR {:.1e}, Loss: {:.1f}'
.format(epoch, Epochs, str(i).zfill(3), str(steps_per_epoch).zfill(3),
scheduler.get_last_lr()[0], \
loss.item()))
model.eval()
mAP_score_regular, mAP_score_ema = validate_multi(val_loader, model, ema, logger, args)
model.train()
if dist.get_rank() == 0:
save_checkpoint({
'state_dict': model.state_dict(),
'epoch': epoch,
'mAP': mAP_score_regular
},
savedir = args.output,
savedname = 'model-{}-{}.ckpt'.format(epoch + 1, i + 1),
is_best = mAP_score_regular > mAP_score_ema and mAP_score_regular > highest_mAP )
save_checkpoint({
'state_dict': ema.module.state_dict(),
'epoch': epoch,
'mAP': mAP_score_ema
},
savedir = args.output,
savedname = 'model-{}-{}-{}.ckpt'.format('ema', epoch + 1, i + 1),
is_best = mAP_score_ema > mAP_score_regular and mAP_score_ema > highest_mAP)
mAP_score = max(mAP_score_regular, mAP_score_ema)
highest_mAP = max(highest_mAP, mAP_score)
logger.info('current_mAP = {:.4f}, highest_mAP = {:.4f}\n'.format(mAP_score, highest_mAP))
def main():
args = parser.parse_args()
args.batch_size = args.batch_size
# setup model
print('creating model...')
state = torch.load(args.model_path, map_location='cuda')
#args.num_classes = state['num_classes']
args.do_bottleneck_head = False
model = create_model(args).cuda()
#load pretrained model
model.load_state_dict(state['model'], strict=True)
in_features = model.head.fc.in_features
#initially, num_class == 1000 in order to load model pretrained on ImageNet correctly, not change it to 80
model.head.fc = torch.nn.Linear(in_features, 80).cuda()
ema = EMA(model, 0.999)
ema.register()
model.train()
classes_list = np.array(list(idx_to_class.values()))
print('done\n')
# Data loading code
normalize = transforms.Normalize(mean=[0, 0, 0], std=[1, 1, 1])
instances_path_val = os.path.join(args.data,
'annotations/instances_val2017.json')
#instances_path_train = os.path.join(args.data, 'annotations/instances_val2017.json')#temprarily use val as train
instances_path_train = os.path.join(
args.data, 'annotations/instances_train2017.json')
data_path_val = os.path.join(args.data, 'val2017')
#data_path_train = os.path.join(args.data, 'val2017')#temporarily use val as train
data_path_train = os.path.join(args.data, 'train2017')
val_dataset = CocoDetection(
data_path_val,
instances_path_val,
transforms.Compose([
transforms.Resize((args.image_size, args.image_size)),
#RandAugment(),
transforms.ToTensor(),
normalize,
]))
train_dataset = CocoDetection(
data_path_train,
instances_path_train,
transforms.Compose([
transforms.Resize((args.image_size, args.image_size)),
RandAugment(),
transforms.ToTensor(),
normalize,
#Cutout(n_holes = 1, length = 16)
Cutout(n_holes=1, length=args.image_size / 2)
]))
print("len(val_dataset)): ", len(val_dataset))
print("len(train_dataset)): ", len(train_dataset))
val_loader = torch.utils.data.DataLoader(val_dataset,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=False)
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=False)
#set optimizer
lr = 0.0002
Epoch = 10
weight_decay = 0.0001
criterion = AsymmetricLoss()
#params = model.parameters()
skip = {}
if hasattr(model, 'no_weight_decay'):
skip = model.no_weight_decay()
parameters = add_weight_decay(model, weight_decay, skip)
weight_decay = 0.
opt_args = dict(lr=lr, weight_decay=weight_decay)
#optimizer = torch.optim.Adam(params, lr=lr)#尝试新的optimizer
optimizer = AdamW(parameters, **opt_args)
total_step = len(train_loader)
scheduler = lr_scheduler.OneCycleLR(optimizer,
max_lr=lr,
total_steps=total_step,
epochs=Epoch)
#total_step = len(train_loader)
highest_mAP = 0
trainInfoList = []
Sig = torch.nn.Sigmoid()
#f=open('info_train.txt', 'a')
for epoch in range(Epoch):
for i, (inputData, target) in enumerate(train_loader):
f = open('info_train.txt', 'a')
#model.train()
inputData = inputData.cuda()
target = target.cuda()
target = target.max(dim=1)[0]
#Sig = torch.nn.Sigmoid()
output = Sig(model(inputData))
#output[output<args.thre] = 0
#output[output>=args.thre]=1
#print(output.shape) #(batchsize, channel, imhsize, imgsize)
#print(inputData.shape) #(batchsize, numclasses)
#print(output[0])
#print(target[0])
loss = criterion(output, target)
model.zero_grad()
loss.backward()
optimizer.step()
ema.update()
#store information
if i % 10 == 0:
trainInfoList.append([epoch, i, loss.item()])
print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'.format(
epoch, Epoch, i, total_step, loss.item()))
f.write('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}\n'.format(
epoch, Epoch, i, total_step, loss.item()))
if (i + 1) % 100 == 0:
#储存相应迭代模型
torch.save(
model.state_dict(),
os.path.join('models/',
'model-{}-{}.ckpt'.format(epoch + 1, i + 1)))
f.close()
scheduler.step() #修改学习率
#validate per epoch and record result
f = open('info_train.txt', 'a')
model.eval()
mAP_score = validate_multi(val_loader, model, args, ema)
model.train()
if mAP_score > highest_mAP:
highest_mAP = mAP_score
print('current highest_mAP = ', highest_mAP)
f.write('current_mAP = {}, highest_mAP = {}\n'.format(
mAP_score, highest_mAP))
torch.save(model.state_dict(),
os.path.join('models/', 'model-highest.ckpt'))
f.close()
def train_multi_label_coco(model, train_loader, val_loader, args):
ema = ModelEma(model, 0.9997) # 0.9997^641=0.82
# set optimizer
Epochs = 80
Stop_epoch = 60
weight_decay = 1e-4
if args.loss.lower() == 'asymmetricloss':
criterion = AsymmetricLoss(gamma_neg=4,
gamma_pos=0,
clip=0.05,
disable_torch_grad_focal_loss=True)
elif args.loss.lower() == 'normalceloss':
criterion = NormalCELoss(reduction='sum')
elif args.loss.lower() == 'bnnceloss':
criterion = BNNCELoss(reduction='sum')
else:
raise ValueError('Wrong loss functon was choosen!')
parameters = add_weight_decay(model, weight_decay)
optimizer = torch.optim.Adam(params=parameters, lr=args.lr,
weight_decay=0) # true wd, filter_bias_and_bn
steps_per_epoch = len(train_loader)
scheduler = lr_scheduler.OneCycleLR(optimizer,
max_lr=args.lr,
steps_per_epoch=steps_per_epoch,
epochs=Epochs,
pct_start=0.2)
highest_mAP = 0
trainInfoList = []
scaler = GradScaler()
writer = SummaryWriter(
comment='_Model_{}_LossFunc_{}_LR_{}_BS_{}_WD_{}'.format(
args.model_name, args.loss, args.lr, args.batch_size,
weight_decay))
for epoch in range(Epochs):
if epoch > Stop_epoch:
break
for i, (inputData,
target) in enumerate(tqdm(train_loader, desc='Training ')):
inputData = inputData.cuda()
target = target.cuda() # (batch,3,num_classes)
target = target.max(dim=1)[0]
with autocast(): # mixed precision
output = model(inputData) # sigmoid will be done in loss !
loss = criterion(output, target, steps_per_epoch, i) # FIXME
model.zero_grad()
scaler.scale(loss).backward()
# loss.backward()
scaler.step(optimizer)
scaler.update()
# optimizer.step()
scheduler.step()
ema.update(model)
writer.add_scalar('loss/train_batch',
loss.item() / target.shape[0],
(epoch * steps_per_epoch) + (i + 1))
trainInfoList.append(loss.item())
writer.add_scalar('loss/train_epoch', sum(trainInfoList), epoch + 1)
trainInfoList = []
model.eval()
mAP_score = validate_multi(val_loader, model, ema)
model.train()
if mAP_score > highest_mAP:
highest_mAP = mAP_score
try:
save_path = os.path.join(
'checkpoints/',
'model-highest-{}.ckpt'.format(args.model_name))
torch.save(model.state_dict(), save_path)
print('model is saved as:', save_path)
except:
pass
print('current_mAP = {:.2f}, highest_mAP = {:.2f}\n'.format(
mAP_score, highest_mAP))
writer.add_scalar('mAP/current', mAP_score, epoch + 1)
writer.add_scalar('mAP/highest', highest_mAP, epoch + 1)