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)