def save_feature_batch(model_name, ckpt_path, feature_path, label_path, data_type="valid"): model = create_model(model_name.split('-')[-2], num_classes=3, checkpoint_path=ckpt_path) model.cuda().eval() print('..... Finished loading model! ......') img_size = int(model_name.split('-')[-1]) interpolation = "bicubic" batch_size = 128 dataset = Dataset(os.path.join(BASE, data_type)) loader = create_loader( dataset, input_size=img_size, batch_size=batch_size, use_prefetcher=False, interpolation=interpolation, num_workers=8) features = [] labels = [] with torch.no_grad(): for batch_idx, (input, target) in enumerate(tqdm(loader)): out = model.forward_features(input.cuda()) out2 = nn.AdaptiveAvgPool2d(1)(out) feature = out2.view(out.size(0), -1) features.append(feature.cpu().numpy()) labels.extend(target.cpu().numpy()) features = np.array(np.vstack(features)) pickle.dump(features, open(feature_path, 'wb')) pickle.dump(labels, open(label_path, 'wb')) print('CNN features obtained and saved.')
def save_feature_test(model_name, ckpt_path, feature_path, imgs): ''' 提取特征,保存为pkl文件 ''' model = create_model(model_name.split('-')[-2], num_classes=3, checkpoint_path=ckpt_path) model.eval().cuda() print('..... Finished loading model! ......') ## 特征的维度需要自己根据特定的模型调整,我这里采用的是哪一个我也忘了 nb_features = 2048 features = np.empty((len(imgs), nb_features)) for i in tqdm(range(len(imgs))): img_path = imgs[i].strip().split(' ')[0] img = Image.open(img_path).convert('RGB') # bilinear bicubic img = get_test_transform(img_size=int(model_name.split('-')[-1]), interpolation="bicubic")(img).unsqueeze( 0).cuda() with torch.no_grad(): out = model.forward_features(img) out2 = nn.AdaptiveAvgPool2d(1)(out) feature = out2.view(out.size(1), -1).squeeze(1) features[i, :] = feature.cpu().numpy() pickle.dump(features, open(feature_path, 'wb')) print('CNN features obtained and saved.')
def main(): args = parser.parse_args() args.batch_size = args.batch_size # setup model print('creating and loading the model...') state = torch.load(args.model_path, map_location='cpu') args.num_classes = state['num_classes'] args.do_bottleneck_head = False model = create_model(args).cuda() model.load_state_dict(state['model'], strict=True) model.eval() classes_list = np.array(list(state['idx_to_class'].values())) print('done\n') # Data loading code normalize = transforms.Normalize(mean=[0, 0, 0], std=[1, 1, 1]) instances_path = os.path.join(args.data, 'annotations/instances_val2017.json') data_path = os.path.join(args.data, 'val2017') val_dataset = CocoDetection(data_path, instances_path, transforms.Compose([ transforms.Resize((args.image_size, args.image_size)), transforms.ToTensor(), normalize, ])) print("len(val_dataset)): ", len(val_dataset)) val_loader = torch.utils.data.DataLoader( val_dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.workers, pin_memory=True) validate_multi(val_loader, model, args)
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.do_bottleneck_head = False # Setup model print('creating model...') model = create_model(args).cuda() if args.model_path: # make sure to load pretrained ImageNet model state = torch.load(args.model_path, map_location='cpu') filtered_dict = {k: v for k, v in state['model'].items() if (k in model.state_dict() and 'head.fc' not in k)} model.load_state_dict(filtered_dict, strict=False) print('done\n') # COCO Data loading instances_path_val = os.path.join(args.data, 'annotations/instances_val2014.json') instances_path_train = os.path.join(args.data, 'annotations/instances_train2014.json') # data_path_val = args.data # data_path_train = args.data data_path_val = f'{args.data}/val2014' # args.data data_path_train = f'{args.data}/train2014' # args.data val_dataset = CocoDetection(data_path_val, instances_path_val, transforms.Compose([ transforms.Resize((args.image_size, args.image_size)), transforms.ToTensor(), # normalize, # no need, toTensor does normalization ])) train_dataset = CocoDetection(data_path_train, instances_path_train, transforms.Compose([ transforms.Resize((args.image_size, args.image_size)), CutoutPIL(cutout_factor=0.5), RandAugment(), transforms.ToTensor(), # normalize, ])) print("len(val_dataset)): ", len(val_dataset)) print("len(train_dataset)): ", len(train_dataset)) # Pytorch Data loader train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=True) val_loader = torch.utils.data.DataLoader( val_dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.workers, pin_memory=False) # Actuall Training train_multi_label_coco(model, train_loader, val_loader, args.lr)
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'] model = create_model(args).cuda() model.load_state_dict(state['model'], strict=True) model.eval() classes_list = np.array(list(state['idx_to_class'].values())) print('done\n') 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 print(tensor_img.shape) tensor_batch = torch.unsqueeze(tensor_img, 0).cuda() print(tensor_batch.shape) output = torch.squeeze(torch.sigmoid(model(tensor_batch))) np_output = output.cpu().detach().numpy() # with open("./np_output.pickle", 'wb') as handle: # pickle.dump(np_output, handle) # with open("./before_sigmod.pickle", 'wb') as handle: # pickle.dump(model(tensor_batch).cpu().detach().numpy(), handle) detected_classes = classes_list[np_output > args.th] print('done\n') print(detected_classes) # 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(): # parsing args args = parser.parse_args() # setup model print('creating model...') model = create_model(args).cuda() state = torch.load(args.model_path, map_location='cpu')['model'] model.load_state_dict(state, strict=False) model.eval() print('done\n') # setup data loader print('creating data loader...') test_loader = create_dataloader(args) print('done\n') # actual validation process print('doing testing...') infer(model, test_loader, args) print('done\n')
def main(): # parsing args args = parser.parse_args() # setup model print('creating model...') model = create_model(args).cuda() state = torch.load(args.model_path, map_location='cpu')['model'] model.load_state_dict(state, strict=True) model.eval() print('done\n') # setup data loader print('creating data loader...') val_loader = create_dataloader(args) print('done\n') # actual validation process print('doing validation...') prec1_f, prec5_f = validate(model, val_loader) print("final top-1 validation accuracy: {:.2f}".format(prec1_f.avg))
def main(): args = parser.parse_args() args.batch_size = args.batch_size # setup logger logger = setup_logger(output=args.output, color=False, name="nuswide") logger.info("Command: "+' '.join(sys.argv)) # setup model print('creating and loading the model...') state = torch.load(args.model_path, map_location='cpu') args.num_classes = state['num_classes'] args.do_bottleneck_head = False model = create_model(args).cuda() model.load_state_dict(state['model'], strict=True) model.eval() print('done\n') # Data loading code normalize = transforms.Normalize(mean=[0, 0, 0], std=[1, 1, 1]) val_dataset = NusWideAslDataset( img_dir=args.img_dir, csv_path=args.csv_path, split='val', transform=transforms.Compose([ transforms.Resize((args.image_size, args.image_size)), transforms.ToTensor(), normalize, ]) ) print("len(val_dataset)): ", len(val_dataset)) val_loader = torch.utils.data.DataLoader( val_dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.workers, pin_memory=True) validate_multi(val_loader, model, args, logger)
def predict(self): dataset = create_dataset( self.opt ) # create a dataset given opt.dataset_mode and other options model = create_model( self.opt) # create a model given opt.model and other options model.setup( self.opt ) # regular setup: load and print networks; create schedulers img_rec = None if self.opt.eval: model.eval() for i, data in enumerate(dataset): model.set_input(data) # unpack data from data loader model.test() # run inference visuals = model.get_current_visuals() # get image results im = visuals['fake_B'] img_rec = util.tensor2im(im) return img_rec
def main(): # parsing args args = parser.parse_args() # gpu setup device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') # setup model print('load pretrained model...') model = create_model(args) state = torch.load(args.model_path, map_location='cpu')['model'] model.load_state_dict(state, strict=False) print('set up the model for fine-grained task...') for param in model.parameters(): param.requires_grad = False # freeze the convnet num_in_fc = model.num_features model.head.fc = nn.Linear(num_in_fc, args.transfer_num_classes) model = model.to(device) print('done\n') # setup data loader print('creating data loader...') train_loader = create_dataloader(args) print('done\n') # training process print('start training...') # only train the fc print('train only fc layer:') model = train(model, train_loader, args, device, stage='only_fc') # train the entire network print('train entire network:') for param in model.parameters(): param.requires_grad = True train(model, train_loader, args, device, stage='full_net') print('finish training\n')
def main(): # parsing args args = parser.parse_args() # setup model print('creating model...') model = create_model(args).cuda() from src.models.tresnet_v2.tresnet_v2 import InplacABN_to_ABN model2 = InplacABN_to_ABN(model) aaa = torch.jit.script(model2) state = torch.load(args.model_path, map_location='cpu')['model'] model.load_state_dict(state, strict=False) model.eval() print('done\n') # setup data loader print('creating data loader...') val_loader = create_dataloader(args) print('done\n') # actual validation process print('doing validation...') prec1_f = validate(model, val_loader) print("final top-1 validation accuracy: {:.2f}".format(prec1_f.avg))
def main(): args = parser.parse_args() args.batch_size = args.batch_size args.do_bottleneck_head = False # setup model print('creating model...') model = create_model(args).cuda() if args.model_path: state = torch.load(args.model_path, map_location='cpu') filtered_dict = { k: v for k, v in state['model'].items() if (k in model.state_dict() and 'head.fc' not in k) } model.load_state_dict(filtered_dict, strict=False) # model.load_state_dict(state['model'], strict=False) 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_val2014.json') instances_path_train = os.path.join( args.data, 'annotations/instances_train2014.json') # data_path_val = os.path.join(args.data, 'val2014') # data_path_train = os.path.join(args.data, 'train2014') data_path_val = args.data data_path_train = args.data val_dataset = CocoDetection( data_path_val, instances_path_val, transforms.Compose([ transforms.Resize((args.image_size, args.image_size)), transforms.ToTensor(), # normalize, # no need, toTensor does normalization ])) train_dataset = CocoDetection( data_path_train, instances_path_train, transforms.Compose([ transforms.Resize((args.image_size, args.image_size)), CutoutPIL(cutout_factor=0.5), RandAugment(), 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) train_multi_label_coco(model, train_loader, val_loader, args.lr, gamma_neg=4, gamma_pos=0, clip=0.05)
def main(): # parsing args args = parser.parse_args() transforms = torchvision.transforms.Compose([ torchvision.transforms.ToPILImage(), torchvision.transforms.Resize((112, 112)), torchvision.transforms.ToTensor(), ]) def load_video(video_path: str): image_size = 112 video = torchvision.io.read_video(video_path, pts_unit="sec")[0].float() # [t, n, h, w] video = video.permute(0, 3, 1, 2) t_size = 16 #new_video = torch.zeros((min(video.shape[0], t_size), video.shape[1], image_size, image_size)) new_video = torch.zeros( (t_size, video.shape[1], image_size, image_size)) for t in range(min(video.shape[0], t_size)): if video.shape[0] < t_size: frame = transforms(video[t]) new_video[t] = (frame - 127.5) / 127.5 else: t_ = int(video.shape[0] / t_size * t) frame = transforms(video[t_]) new_video[t] = (frame - 127.5) / 127.5 return new_video def validate(model, val_loader): device = 'cuda' model.eval() total, correct = 0, 0 for j, (video, label) in enumerate(val_loader): #batch = torch.tensor(batch) #video, label = batch[:,0], batch[:,1] video, label = video.view(-1, 3, 112, 112), label.view(-1) video = video.to(device) #label = label.cuda() with torch.no_grad(): pred = model(video) pred = F.softmax(pred, dim=1) pred = torch.argmax(pred, dim=1).cpu().numpy() correct += accuracy_score(label, pred, normalize=False) total += 16 #batch_size return correct / total * 100 def train(model, train_loader, val_loader, loss_fn, optimizer, scheduler, max_epochs): device = 'cuda' model.train() loss_fn = loss_fn.to(device) optimizer.zero_grad() #loss = torch.as_tensor(0.0).cuda() best_acc = 0 best_model = copy.deepcopy(model) for i in range(max_epochs): for j, (video, label) in enumerate(train_loader): #batch = torch.tensor(batch) #video, label = batch[:,0], batch[:,1] video, label = video.view(-1, 3, 112, 112), label.view(-1) video = video.to(device) label = label.to(device) pred = model(video) pred = F.softmax(pred, dim=1) loss = loss_fn(pred, label) #torch.cuda.empty_cache() #loss /= len(batch) loss.backward() optimizer.step() scheduler.step() if j % 100 == 0: print(f"epoch:{i}, {j}, loss:{loss.item()}") wandb.log({"loss": loss.item()}) if j % 1000 == 0: acc = validate(model, val_loader) wandb.log({"acc": acc}) if acc > best_acc: best_acc = acc best_model = copy.deepcopy(model) print("best_acc:", best_acc) return best_model wandb.init(project="stam", name=f"stam-{datetime.datetime.now()}") # k times cross validation for i in range(1, 4): cfg = Config() device = "cuda" train_dataset = UCF101( "./dataset/UCF101", [ f"./dataset/ucfTrainTestlist/trainlist0{i}.txt", f"./dataset/ucfTrainTestlist/trainlist0{((i+1)%3)+1}.txt", ], ) testset = UCF101( "./dataset/UCF101", [f"./dataset/ucfTrainTestlist/trainlist0{((i+2)%3)+1}.txt"], ) t = len(testset) t, v = int(t * 0.9), t - int(t * 0.9) test_dataset, val_dataset = torch.utils.data.random_split( testset, [t, v]) train_loader = DataLoader( train_dataset, 1, shuffle=True, #collate_fn=video_collate, ) val_loader = DataLoader( val_dataset, 1, shuffle=True, #collate_fn=video_collate, ) test_loader = DataLoader( test_dataset, 1, shuffle=False, #collate_fn=video_collate, ) # setup model print('creating model...') model = create_model(args).to(device) #state = torch.load(args.model_path, map_location='cpu')['model'] #model.load_state_dict(state, strict=False) #model.eval() print('done\n') optimizer = torch.optim.Adam(model.parameters(), lr=cfg.base_lr) scheduler = torch.optim.lr_scheduler.LambdaLR( optimizer, lambda x: (512**-0.5) * min((x + 1)**(-0.5), (x + 1) * cfg.warmup_steps**(-1.5)), ) loss_fn = nn.CrossEntropyLoss() print(sum(p.numel() for p in model.parameters())) model = train(model, train_loader, val_loader, loss_fn, optimizer, scheduler, max_epochs=cfg.epochs) torch.save(model, f"./checkpoints/ckpt-{i}.pt") acc = validate(model, test_loader) print(f"{i}acc:{acc}") wandb.log({"test_acc": acc})
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() #修改学习率
parser.add_argument('--model_path') parser.add_argument('--model_name', type=str, default='tresnet_m') parser.add_argument('--num_classes', type=int, default=1000) parser.add_argument('--input_size', type=int, default=224) parser.add_argument('--val_zoom_factor', type=int, default=0.875) parser.add_argument('--batch_size', type=int, default=48) parser.add_argument('--num_workers', type=int, default=8) # parsing args args = parser.parse_args() val_tfms = create_val_tfm(args) #### TResNet-M #### args.model_name = 'tresnet_m' model_path = './tresnet_m.pth' model = create_model(args) state = torch.load(model_path, map_location='cpu')['model'] model.load_state_dict(state, strict=True) model = InplacABN_to_ABN(model) model = fuse_bn_recursively(model) model = model.cuda() model = model.half() model.eval() # Run the benchmark print('Benchmarking TResNet-M') for i in range(1): # Two times for caching ImageNet.benchmark( model=model, paper_model_name='TResNet-M-FP16', paper_arxiv_id='2003.13630',
def validate(args): # might as well try to validate something args.prefetcher = not args.no_prefetcher # create model model = create_model( args.model, num_classes=args.num_classes, in_chans=3, global_pool=args.gp) if args.checkpoint: load_checkpoint(model, args.checkpoint) param_count = sum([m.numel() for m in model.parameters()]) _logger.info('Model %s created, param count: %d' % (args.model, param_count)) data_config = resolve_data_config(vars(args), model=model) # model, test_time_pool = apply_test_time_pool(model, data_config, args) if torch.cuda.is_available(): model.cuda() criterion = nn.CrossEntropyLoss().cuda() dataset = Dataset(args.data) crop_pct = data_config['crop_pct'] loader = create_loader( dataset, input_size=data_config['input_size'], batch_size=args.batch_size, use_prefetcher=args.prefetcher, interpolation=data_config['interpolation'], mean=data_config['mean'], std=data_config['std'], num_workers=args.workers, crop_pct=crop_pct) batch_time = AverageMeter() losses = AverageMeter() top1 = AverageMeter() top5 = AverageMeter() f1_m = AverageMeter() end = time.time() total_pred_idx = [] total_truth_idx = [] mistake_image = [] mistake_image_dict = {'calling': [], 'normal': [], 'smoking': [], 'smoking_calling': []} # class_2_index = {0: 'normal', 1: 'phone', 2: 'smoke'} class_2_index = {0: 'calling', 1: 'normal', 2: 'smoking', 3: 'smoking_calling'} with open("./txts/%s.json" % json_name, 'r', encoding="utf-8") as f: shape_dict = json.load(f) dets_info = {} model.eval() with torch.no_grad(): # warmup, reduce variability of first batch time, especially for comparing torchscript vs non input = torch.randn((args.batch_size,) + data_config['input_size']) if torch.cuda.is_available(): input = input.cuda() model(input) end = time.time() for batch_idx, (input, target) in enumerate(loader): if args.no_prefetcher and torch.cuda.is_available(): target = target.cuda() input = input.cuda() # compute output # t0 = time.time() output = model(input) # print("time0: %.8f s" % ((time.time() - t0))) # t1 = time.time() # out = output.detach().cpu() # print("time1: %.8f s" % ((time.time() - t1) / 64)) # print("time2: %.8f s" % ((time.time() - t0) / 64)) # t2 = time.time() # out = out.cuda().cpu() # print("time3: %.8f s" % ((time.time() - t2) / 64)) # get prediction index and ground turth index prob = torch.max(F.softmax(output, -1), -1)[0] idx = torch.max(F.softmax(output, -1), -1)[1] target_idx = target.cpu().numpy() predict_idx = idx.cpu().numpy() for j in range(len(target_idx)): total_truth_idx.append(target_idx[j]) total_pred_idx.append(predict_idx[j]) class_dict = loader.dataset.class_to_idx target_class = list(class_dict.keys())[list(class_dict.values()).index(int(target_idx[j]))] pred_class = list(class_dict.keys())[list(class_dict.values()).index(int(predict_idx[j]))] filename = loader.dataset.filenames()[batch_idx * args.batch_size + j] name = filename.split('/')[-1].split('.')[0] dets_info[name] = [pred_class, float(prob[j]), shape_dict[name][1], shape_dict[name][2]] if target_idx[j] != predict_idx[j]: mistake_image.append( [loader.dataset.filenames()[batch_idx * args.batch_size + j], target_class, pred_class, np.round(prob[j].cpu().numpy(), 4)]) mistake_image_dict[class_2_index[predict_idx[j]]].append( loader.dataset.filenames()[batch_idx * args.batch_size + j]) loss = criterion(output, target) # measure accuracy and record loss prec1, prec5 = accuracy(output.data, target, topk=(1, 3)) losses.update(loss.item(), input.size(0)) top1.update(prec1.item(), input.size(0)) top5.update(prec5.item(), input.size(0)) # measure elapsed time batch_time.update(time.time() - end) end = time.time() if batch_idx % args.log_freq == 0: _logger.info( 'Test: [{0:>4d}/{1}] ' 'Time: {batch_time.val:.3f}s ({batch_time.avg:.3f}s, {rate_avg:>7.2f}/s) ' 'Loss: {loss.val:>7.4f} ({loss.avg:>6.4f}) ' 'Acc@1: {top1.val:>7.2f} ({top1.avg:>7.2f}) ' 'Acc@5: {top5.val:>7.3f} ({top5.avg:>7.3f})'.format( batch_idx, len(loader), batch_time=batch_time, rate_avg=input.size(0) / batch_time.avg, loss=losses, top1=top1, top5=top5)) with open("%s/%s.json" % (output_path, json_name.split('-')[0]), "w", encoding="utf-8") as f: json.dump(dets_info, f) top1a, top5a = top1.avg, top5.avg results = OrderedDict( top1=round(top1a, 4), top1_err=round(100 - top1a, 4), top5=round(top5a, 4), top5_err=round(100 - top5a, 4), param_count=round(param_count / 1e6, 2), img_size=data_config['input_size'][-1], cropt_pct=crop_pct, interpolation=data_config['interpolation'], mistake_image_dict=mistake_image_dict, pred_idx=total_pred_idx, truth_idx=total_truth_idx) _logger.info(' * Acc@1 {:.2f} ({:.2f}) Acc@5 {:.2f} ({:.2f})'.format( results['top1'], results['top1_err'], results['top5'], results['top5_err'])) map, each_ap = eval_map(detFolder="%s/%s.json" % (output_path, json_name.split('-')[0]), gtFolder="txts/%s.json" % json_name, return_each_ap=True) _logger.info('Valid mAP: {}, each ap: {}'.format(round(map, 4), each_ap)) return results
def model_train(): """ Trains model which is used as a feature extractor :return:None """ # Download data downloadData(data_path='/input/speech_commands/') # Get data dictionary dataDict = getDataDict(data_path='/input/speech_commands/') # Obtain dataframe for each dataset trainDF = getDataframe(dataDict['train']) valDF = getDataframe(dataDict['val']) devDF = getDataframe(dataDict['dev']) testDF = getDataframe(dataDict['test']) print("Dataset statistics") print("Train files: {}".format(trainDF.shape[0])) print("Validation files: {}".format(valDF.shape[0])) print("Dev test files: {}".format(devDF.shape[0])) print("Test files: {}".format(testDF.shape[0])) # Use TF Data API for efficient data input train_data, train_steps = getDataset(df=trainDF, batch_size=BATCH_SIZE, cache_file='train_cache', shuffle=True) val_data, val_steps = getDataset(df=valDF, batch_size=BATCH_SIZE, cache_file='val_cache', shuffle=False) model = create_model() model.summary() # Stop training if the validation accuracy doesn't improve earlyStopping = EarlyStopping(monitor='val_loss', patience=PATIENCE, verbose=1) # Reduce LR on validation loss plateau reduceLR = ReduceLROnPlateau(monitor='val_loss', patience=PATIENCE, verbose=1) # Compile the model model.compile(loss='sparse_categorical_crossentropy', optimizer=Adam(learning_rate=LEARNING_RATE), metrics=["sparse_categorical_accuracy"]) # Train the model history = model.fit(train_data.repeat(), steps_per_epoch=train_steps, validation_data=val_data.repeat(), validation_steps=val_steps, epochs=EPOCHS, callbacks=[earlyStopping, reduceLR]) # Save model print("Saving model") model.save('../models/marvin_kws.h5') # Save history data print("Saving training history") with open('../models/marvin_kws_history.pickle', "wb") as file: pickle.dump(history.history, file, protocol=pickle.HIGHEST_PROTOCOL) plot_history(history=history)
def main(): setup_default_logging() args, args_text = _parse_args() args.prefetcher = not args.no_prefetcher torch.manual_seed(args.seed) model = create_model(args.model, pretrained=True, num_classes=args.num_classes, drop_rate=args.drop, drop_path_rate=args.drop_path, drop_block_rate=args.drop_block, checkpoint_path=args.initial_checkpoint) if args.local_rank == 0: _logger.info('Model %s created, param count: %d' % (args.model, sum([m.numel() for m in model.parameters()]))) data_config = resolve_data_config(vars(args), model=model, verbose=args.local_rank == 0) if args.num_gpu > 1: model = nn.DataParallel(model, device_ids=list(range(args.num_gpu))).cuda() else: model.cuda() optimizer = create_optimizer(args, model) loss_scaler = None # optionally resume from a checkpoint resume_epoch = None if args.resume: resume_epoch = resume_checkpoint( model, args.resume, optimizer=None if args.no_resume_opt else optimizer, loss_scaler=None if args.no_resume_opt else loss_scaler, log_info=args.local_rank == 0) lr_scheduler, num_epochs = create_scheduler(args, optimizer) start_epoch = 0 if args.start_epoch is not None: # a specified start_epoch will always override the resume epoch start_epoch = args.start_epoch elif resume_epoch is not None: start_epoch = resume_epoch if lr_scheduler is not None and start_epoch > 0: lr_scheduler.step(start_epoch) if args.local_rank == 0: _logger.info('Scheduled epochs: {}'.format(num_epochs)) train_dir = os.path.join(args.data, 'train') if not os.path.exists(train_dir): _logger.error( 'Training folder does not exist at: {}'.format(train_dir)) exit(1) dataset_train = Dataset(train_dir) collate_fn = None mixup_fn = None mixup_active = args.mixup > 0 or args.cutmix > 0. or args.cutmix_minmax is not None if mixup_active: mixup_args = dict(mixup_alpha=args.mixup, cutmix_alpha=args.cutmix, cutmix_minmax=args.cutmix_minmax, prob=args.mixup_prob, switch_prob=args.mixup_switch_prob, elementwise=args.mixup_elem, label_smoothing=args.smoothing, num_classes=args.num_classes) if args.prefetcher: collate_fn = FastCollateMixup(**mixup_args) else: mixup_fn = Mixup(**mixup_args) loader_train = create_loader( dataset_train, input_size=data_config['input_size'], batch_size=args.batch_size, is_training=True, use_prefetcher=args.prefetcher, re_prob=args.reprob, re_mode=args.remode, re_count=args.recount, color_jitter=args.color_jitter, auto_augment=args.aa, num_workers=args.workers, collate_fn=collate_fn, ) eval_dir = os.path.join(args.data, 'valid') if not os.path.isdir(eval_dir): eval_dir = os.path.join(args.data, 'validation') if not os.path.isdir(eval_dir): _logger.error( 'Validation folder does not exist at: {}'.format(eval_dir)) exit(1) dataset_eval = Dataset(eval_dir) loader_eval = create_loader( dataset_eval, input_size=data_config['input_size'], batch_size=args.batch_size, is_training=False, use_prefetcher=args.prefetcher, num_workers=args.workers, crop_pct=data_config['crop_pct'], ) if mixup_active: # smoothing is handled with mixup target transform train_loss_fn = SoftTargetCrossEntropy().cuda() elif args.smoothing: train_loss_fn = LabelSmoothingCrossEntropy( smoothing=args.smoothing).cuda() train_loss_ce = nn.CrossEntropyLoss().cuda() else: train_loss_fn = nn.CrossEntropyLoss().cuda() validate_loss_fn = nn.CrossEntropyLoss().cuda() eval_metric = args.eval_metric best_metric = None best_epoch = None saver = None output_dir = '' plateau_num = 0 if args.local_rank == 0: output_base = args.output if args.output else './output' exp_name = '-'.join([ datetime.now().strftime("%Y%m%d-%H%M%S"), args.model, str(data_config['input_size'][-1]) ]) output_dir = get_outdir(output_base, exp_name) decreasing = True if eval_metric == 'loss' else False saver = CheckpointSaver(model=model, optimizer=optimizer, args=args, amp_scaler=loss_scaler, checkpoint_dir=output_dir, recovery_dir=output_dir, decreasing=decreasing, max_history=2) with open(os.path.join(output_dir, 'args.yaml'), 'w') as f: f.write(args_text) with open("./txts/%s.json" % json_name, 'r', encoding="utf-8") as f: shape_dict = json.load(f) try: for epoch in range(start_epoch, num_epochs): train_metrics = train_epoch(epoch, model, loader_train, optimizer, [train_loss_fn, train_loss_ce], args, lr_scheduler=lr_scheduler, output_dir=output_dir, mixup_fn=mixup_fn) eval_metrics, dets_info = validate(model, loader_eval, validate_loss_fn, args, shape_dict=shape_dict) with open("%s/v.json" % output_dir, "w", encoding="utf-8") as f: json.dump(dets_info, f) map = round( eval_map(detFolder="%s/v.json" % output_dir, gtFolder="txts/%s.json" % json_name), 4) eval_metrics["map"] = map _logger.info('Valid mAP: {}'.format(map)) if lr_scheduler is not None: # step LR for next epoch lr_scheduler.step(epoch + 1, eval_metrics[eval_metric]) update_summary(epoch, train_metrics, eval_metrics, os.path.join(output_dir, 'summary.csv'), write_header=best_metric is None) if saver is not None: # save proper checkpoint with eval metric save_metric = eval_metrics[eval_metric] saver.save_prefix = "%.2f-%s" % (eval_metrics["top1"], map) best_metric, best_epoch = saver.save_checkpoint( epoch, metric=save_metric) if eval_metrics[eval_metric] >= best_metric: plateau_num = 0 else: plateau_num += 1 # 超过30个epoch还没有更新metric,停止运行 if plateau_num == 30: break except KeyboardInterrupt: pass if best_metric is not None: _logger.info('*** Best metric: {0} (epoch {1})'.format( best_metric, best_epoch))
def create_ensembles(shape, l_name, contamination=0.01): """ Utility function for creating the ensembles :param shape: The input shape :param l_name: The idenfitier of the local outlier detector :param contamination: The contamination (for some models this is a parameter) :return: array(C), array(L) """ num_clients = shape[0] c = create_models(num_clients, shape[-1], compression_factor=0.4) l = None if l_name == "lof1": l = [ LocalOutlierFactor(n_neighbors=1, contamination=contamination, novelty=True) for _ in range(num_clients) ] if l_name == "lof2": l = [ LocalOutlierFactor(n_neighbors=2, contamination=contamination, novelty=True) for _ in range(num_clients) ] if l_name == "lof4": l = [ LocalOutlierFactor(n_neighbors=4, contamination=contamination, novelty=True) for _ in range(num_clients) ] if l_name == "lof8": l = [ LocalOutlierFactor(n_neighbors=8, contamination=contamination, novelty=True) for _ in range(num_clients) ] if l_name == "lof16": l = [ LocalOutlierFactor(n_neighbors=16, contamination=contamination, novelty=True) for _ in range(num_clients) ] if l_name == "lof32": l = [ LocalOutlierFactor(n_neighbors=32, contamination=contamination, novelty=True) for _ in range(num_clients) ] if l_name == "lof64": l = [ LocalOutlierFactor(n_neighbors=64, contamination=contamination, novelty=True) for _ in range(num_clients) ] if l_name == "lof100": l = [ LocalOutlierFactor(n_neighbors=100, contamination=contamination, novelty=True) for _ in range(num_clients) ] if l_name == "xstream": l = [Chains(k=100, nchains=100, depth=15) for _ in range(num_clients)] if l_name == "ae": l = [ create_model(shape[-1], compression_factor=0.4) for _ in range(num_clients) ] if l_name == "if": l = [ IsolationForest(contamination=contamination) for _ in range(num_clients) ] if not l: raise KeyError("No valid local outlier detector name provided.") return np.array(c), np.array(l)
else: data_path = "/home/data/classification/action/new_data/test/" if data_path.split('/')[-1] == 'valid': dataset = Dataset(data_path, transform=tta_test_transform(img_size)) save_weights = True else: dataset = TestDataset(data_path, transform=tta_test_transform(img_size)) save_weights = False checkpoint = glob.glob( os.path.join(output_root, checkpoint_list[index] + '/*best*.pth.tar'))[0] model = create_model('%s' % model_name.split('-')[-2], num_classes=args.num_classes, in_chans=3, checkpoint_path='%s' % checkpoint) model = model.cuda() model.eval() test_loader = torch.utils.data.DataLoader(dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.workers, pin_memory=True) model.default_cfg['model_name'] = model_list[index] prediction_output, true_output = perform_predict( model, test_loader, model_weight,
def main(): args = parser.parse_args() args.do_bottleneck_head = False # setup dist training if 'WORLD_SIZE' in os.environ: args.distributed = int(os.environ['WORLD_SIZE']) > 1 args.device = 'cuda:0' args.world_size = 1 args.rank = 0 # global rank if args.distributed: args.device = 'cuda:%d' % args.local_rank torch.cuda.set_device(args.local_rank) torch.distributed.init_process_group(backend='nccl', init_method='env://') args.world_size = torch.distributed.get_world_size() args.rank = torch.distributed.get_rank() print('Training in distributed mode with multiple processes, 1 GPU per process. Process %d, total %d.' % (args.rank, args.world_size)) else: print('Training with a single process on 1 GPUs.') assert args.rank >= 0 # setup logger logger = setup_logger(output=args.output, distributed_rank=dist.get_rank(), color=False, name="Coco") logger.info("Command: "+' '.join(sys.argv)) if dist.get_rank() == 0: path = os.path.join(args.output, "config.json") with open(path, 'w') as f: json.dump(vars(args), f, indent=2) logger.info("Full config saved to {}".format(path)) os.makedirs(osp.join(args.output, 'tmpdata'), exist_ok=True) # Setup model logger.info('creating model...') model = create_model(args).cuda() if args.model_path: # make sure to load pretrained ImageNet model state = torch.load(args.model_path, map_location='cpu') filtered_dict = {k: v for k, v in state['model'].items() if (k in model.state_dict() and 'head.fc' not in k)} model.load_state_dict(filtered_dict, strict=False) logger.info('done\n') ema = ModelEma(model, 0.9997) # 0.9997^641=0.82 if args.distributed: model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank], broadcast_buffers=False) # COCO Data loading instances_path_val = os.path.join(args.data, 'annotations/instances_val2014.json') instances_path_train = os.path.join(args.data, 'annotations/instances_train2014.json') data_path_val = args.data data_path_train = args.data val_dataset = CocoDetection(data_path_val, instances_path_val, transforms.Compose([ transforms.Resize((args.image_size, args.image_size)), transforms.ToTensor(), # normalize, # no need, toTensor does normalization ])) train_dataset = CocoDetection(data_path_train, instances_path_train, transforms.Compose([ transforms.Resize((args.image_size, args.image_size)), CutoutPIL(cutout_factor=0.5), RandAugment(), transforms.ToTensor(), # normalize, ])) logger.info("len(val_dataset)): {}".format(len(val_dataset))) logger.info("len(train_dataset)): {}".format(len(train_dataset))) train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset) val_sampler = torch.utils.data.distributed.DistributedSampler(val_dataset, shuffle=False) assert args.batch_size // dist.get_world_size() == args.batch_size / dist.get_world_size(), 'Batch size is not divisible by num of gpus.' # Pytorch Data loader train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=args.batch_size // dist.get_world_size(), shuffle=not args.distributed, num_workers=args.workers, pin_memory=True, sampler=train_sampler, drop_last=True) val_loader = torch.utils.data.DataLoader( val_dataset, batch_size=args.batch_size // dist.get_world_size(), shuffle=False, num_workers=args.workers, pin_memory=False, sampler=val_sampler) # Actuall Training train_multi_label_coco(model, ema, train_loader, val_loader, args.lr, args, logger)
def main(): setup_default_logging() args = parser.parse_args() # might as well try to do something useful... args.pretrained = args.pretrained or not args.checkpoint if not os.path.exists(args.output_dir): os.makedirs(args.output_dir) args.checkpoint = glob.glob(args.checkpoint + '/*.pth')[0] # create model model = create_model(args.model, num_classes=args.num_classes, in_chans=3, pretrained=args.pretrained) load_checkpoint(model, args.checkpoint) logging.info('Model %s created, param count: %d' % (args.model, sum([m.numel() for m in model.parameters()]))) args.img_size = int(args.checkpoint.split('/')[-2].split('-')[-1]) config = resolve_data_config(vars(args), model=model) # model, test_time_pool = apply_test_time_pool(model, config, args) if torch.cuda.is_available(): model = model.cuda() loader = create_loader(Dataset(args.data), input_size=config['input_size'], batch_size=args.batch_size, use_prefetcher=False, interpolation=config['interpolation'], mean=config['mean'], std=config['std'], num_workers=args.workers, crop_pct=config['crop_pct']) model.eval() batch_time = AverageMeter() end = time.time() topk_ids = [] scores = [] total_pred_idx = [] total_truth_idx = [] with torch.no_grad(): for batch_idx, (input, target) in enumerate(loader): if torch.cuda.is_available(): input = input.cuda() output = model(input) prob = torch.max(F.softmax(output, -1), -1)[0] idx = torch.max(F.softmax(output, -1), -1)[1] total_pred_idx.extend(idx.cpu().numpy()) total_truth_idx.extend(target.cpu().numpy()) scores.extend(prob.cpu().numpy()) topk_ids.extend(idx.cpu().numpy()) # measure elapsed time batch_time.update(time.time() - end) end = time.time() if batch_idx % args.log_freq == 0: logging.info( 'Predict: [{0}/{1}] Time {batch_time.val:.3f} ({batch_time.avg:.3f})' .format(batch_idx, len(loader), batch_time=batch_time)) # result_file_path = os.path.join('./results', model_name) # if not os.path.exists(result_file_path): # os.makedirs(result_file_path) # res_cf = open('%s/results-all.csv' % result_file_path, mode='w') # for i in range(len(total_pred_idx)): # res_cf.write('{0},'.format(str(total_pred_idx[i]))) # res_cf.write('\n') # for i in range(len(total_truth_idx)): # res_cf.write('{0},'.format(str(total_truth_idx[i]))) # dst_root = './infer/%s' % args.checkpoint.split('/')[-2] # if not os.path.exists(dst_root): # os.makedirs(dst_root) # else: # shutil.rmtree(dst_root) result_list = [] # class_2_index = {0: 'normal', 1: 'calling', 2: 'smoking'} class_2_index = { 0: 'calling', 1: 'normal', 2: 'smoking', 3: 'smoking_calling' } with open(os.path.join(args.output_dir, 'result.json'), 'w', encoding="utf-8") as out_file: filenames = loader.dataset.filenames() for i in range(len(scores)): filename = filenames[i].split('/')[-1] name = class_2_index[topk_ids[i]] result_data = { "image_name": str(filename), "category": name, "score": scores[i] } result_list.append(result_data) # if scores[i] > 0.95: # dst_path = os.path.join(dst_root, name) # if not os.path.exists(dst_path): # os.makedirs(dst_path) # shutil.copy(filenames[i], os.path.join(dst_path, filename)) json.dump(result_list, out_file, cls=MyEncoder, indent=4)
from src.data import create_dataset from src.models import create_model from src.util.visualizer import save_images_to_path import sys if __name__ == '__main__': opt = TestOptions().parse() # get test options # hard-code some parameters for test opt.num_threads = 0 # test code only supports num_threads = 1 opt.batch_size = 1 # test code only supports batch_size = 1 opt.serial_batches = True # disable data shuffling; comment this line if results on randomly chosen images are needed. opt.no_flip = True # no flip; comment this line if results on flipped images are needed. opt.display_id = -1 # no visdom display; the test code saves the results to a HTML file. dataset = create_dataset( opt) # create a dataset given opt.dataset_mode and other options model = create_model( opt) # create a model given opt.model and other options model.setup( opt) # regular setup: load and print networks; create schedulers if opt.force_test_output == '': print('--force_test_output must be defined.') sys.exit() if opt.eval: model.eval() for i, data in enumerate(dataset): if i >= opt.num_test: # only apply our model to opt.num_test images. break model.set_input(data) # unpack data from data loader model.test() # run inference visuals = model.get_current_visuals() # get image results
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()