def main(): config = yaml.load(open("./config/config.yaml", "r"), Loader=yaml.FullLoader) device = "cuda" if torch.cuda.is_available() else "cpu" print(f"Device: {device}") data_transforms = get_simclr_data_transforms(**config["data_transforms"]) train_dataset = datasets.STL10( "./data/", split="train+unlabeled", download=True, transform=MultiViewDataInjector([data_transforms, data_transforms]), ) # online network online_network = ResNet(**config["network"]).to(device) pretrained_folder = config["network"]["fine_tune_from"] # load pre-trained model if pretrained_folder: try: checkpoints_folder = os.path.join("./runs", pretrained_folder, "checkpoints") load_params = torch.load( os.path.join(os.path.join(checkpoints_folder, "model.pth")), map_location=torch.device(device), ) online_network.load_state_dict( load_params["online_network_state_dict"]) except FileNotFoundError: print( "Pre-trained weights not found, starting training from scratch." ) predictor = MlpHead( in_channels=online_network.projection.head[-1].out_features, **config["network"]["projection_head"], ).to(device) # target encoder target_network = ResNet(**config["network"]).to(device) optimizer = torch.optim.SGD( list(online_network.parameters()) + list(predictor.parameters()), **config["optimizer"]["params"], ) trainer = BYOLTrainer( online_network=online_network, target_network=target_network, optimizer=optimizer, predictor=predictor, device=device, **config["trainer"], ) trainer.train(train_dataset)
def main(): config = yaml.load(open("./config/config.yaml", "r"), Loader=yaml.FullLoader) device = 'cuda' if torch.cuda.is_available() else 'cpu' print(f"Training with: {device}") data_transform = get_simclr_data_transforms(**config['data_transforms']) train_dataset = datasets.STL10('/home/thalles/Downloads/', split='train+unlabeled', download=True, transform=MultiViewDataInjector( [data_transform, data_transform])) # online network online_network = ResNet18(**config['network']).to(device) pretrained_folder = config['network']['fine_tune_from'] # load pre-trained model if defined if pretrained_folder: try: checkpoints_folder = os.path.join('./runs', pretrained_folder, 'checkpoints') # load pre-trained parameters load_params = torch.load( os.path.join(os.path.join(checkpoints_folder, 'model.pth')), map_location=torch.device(torch.device(device))) if 'online_network_state_dict' in load_params: online_network.load_state_dict( load_params['online_network_state_dict']) print("Parameters successfully loaded.") except FileNotFoundError: print("Pre-trained weights not found. Training from scratch.") # predictor network predictor = MLPHead( in_channels=online_network.projetion.net[-1].out_features, **config['network']['projection_head']).to(device) # target encoder target_network = ResNet18(**config['network']).to(device) optimizer = torch.optim.SGD( list(online_network.parameters()) + list(predictor.parameters()), **config['optimizer']['params']) trainer = BYOLTrainer(online_network=online_network, target_network=target_network, optimizer=optimizer, predictor=predictor, device=device, **config['trainer']) trainer.train(train_dataset)
def main(): parser = ArgumentParser() parser.add_argument('--incr', action='store_true', help='train representation incrementally') parser.add_argument('--id', type=str, default='', dest='experiment_id', help='experiment id appended to saved files') args = parser.parse_args() config = yaml.load(open("./config/config.yaml", "r"), Loader=yaml.FullLoader) n_class_epochs = config['other']['n_class_epochs'] eval_step = config['other']['eval_step'] device = 'cuda' if torch.cuda.is_available() else 'cpu' print(f"Training with: {device}") data_transform = get_simclr_data_transforms(**config['data_transforms']) #train_dataset = datasets.STL10('../../data', split='train+unlabeled', download=False, # transform=MultiViewDataInjector([data_transform, data_transform])) # online network online_network = ResNet18(**config['network']).to(device) pretrained_folder = config['network']['fine_tune_from'] # load pre-trained model if defined if pretrained_folder: try: checkpoints_folder = os.path.join('./runs', pretrained_folder, 'checkpoints') # load pre-trained parameters load_params = torch.load(os.path.join(os.path.join(checkpoints_folder, 'model.pth')), map_location=torch.device(torch.device(device))) online_network.load_state_dict(load_params['online_network_state_dict']) except FileNotFoundError: print("Pre-trained weights not found. Training from scratch.") # predictor network predictor = MLPHead(in_channels=online_network.projection.net[-1].out_features, **config['network']['projection_head']).to(device) # target encoder target_network = ResNet18(**config['network']).to(device) optimizer = torch.optim.SGD(list(online_network.parameters()) + list(predictor.parameters()), **config['optimizer']['params']) trainer = BYOLTrainer(online_network=online_network, target_network=target_network, optimizer=optimizer, predictor=predictor, device=device, **config['trainer']) num_workers = config['trainer']['num_workers'] batch_size_train = 100 batch_size_test = 200 if args.incr: incr_train_loaders, incr_val_loaders = get_dataloader_incr(data_dir='../../data', base='CIFAR10', num_classes=10, img_size=224, classes_per_exposure=2, train=True, num_workers=num_workers, batch_size_train=batch_size_train, batch_size_test=batch_size_test, transform=MultiViewDataInjector([data_transform, data_transform])) # get train and val indices sampled train_indices = np.concatenate([ldr.sampler.indices for ldr in incr_train_loaders]) train_class_dataloader = DataLoader(incr_train_loaders[0].dataset, sampler=SubsetRandomSampler(train_indices), batch_size=batch_size_train, num_workers=num_workers) #trainer.train(train_dataset) trainer.train_incr(incr_train_loaders, incr_val_loaders, n_class_epochs=n_class_epochs, train_class_dataloader=train_class_dataloader, experiment_id=args.experiment_id, eval_step=eval_step) else: train_loader, val_loader = get_dataloader(data_dir='../../data', base='CIFAR10', num_classes=10, img_size=224, train=True, num_workers=num_workers, batch_size_train=batch_size_train, batch_size_test=batch_size_test, transform=MultiViewDataInjector([data_transform, data_transform])) trainer.train(train_loader, val_loader, n_class_epochs=n_class_epochs, experiment_id=args.experiment_id, eval_step=eval_step)
def main(): config = yaml.load(open("./config/config.yaml", "r"), Loader=yaml.FullLoader) device = 'cuda' if torch.cuda.is_available() else 'cpu' print(f"Training with: {device}") data_transform = get_simclr_data_transforms(**config['data_transforms']) data_transform2 = get_simclr_data_transforms(**config['data_transforms'], blur=1.) # data_transform = get_simclr_data_transforms_randAugment(config['data_transforms']['input_shape']) # data_transform2 = get_simclr_data_transforms_randAugment(config['data_transforms']['input_shape']) train_dataset = datasets.STL10('/media/snowflake/Data/', split='train+unlabeled', download=True, transform=MultiViewDataInjector( [data_transform, data_transform2])) # train_dataset = STL("/home/snowflake/Descargas/STL_data/unlabeled_images", 96, # transform=MultiViewDataInjector([data_transform, data_transform2])) # online network (the one that is trained) online_network = ResNet(**config['network']).to(device) # online_network = MLPmixer(**config['network']).to(device) # target encoder target_network = ResNet(**config['network']).to(device) # target_network = MLPmixer(**config['network']).to(device) pretrained_folder = config['network']['fine_tune_from'] # load pre-trained model if defined if pretrained_folder: try: checkpoints_folder = os.path.join('./runs', pretrained_folder, 'checkpoints') # load pre-trained parameters load_params = torch.load( os.path.join(os.path.join(checkpoints_folder, 'model.pth')), map_location=torch.device(torch.device(device))) online_network.load_state_dict( load_params['online_network_state_dict']) target_network.load_state_dict( load_params['target_network_state_dict']) except FileNotFoundError: print("Pre-trained weights not found. Training from scratch.") # predictor network predictor = MLPHead( in_channels=online_network.projetion.net[-1].out_features, **config['network']['projection_head']).to(device) optimizer = torch.optim.SGD( list(online_network.parameters()) + list(predictor.parameters()), **config['optimizer']['params']) trainer = BYOLTrainer(online_network=online_network, target_network=target_network, optimizer=optimizer, predictor=predictor, device=device, **config['trainer']) trainer.train(train_dataset)
def main(args): config = yaml.load(open("./config/config.yaml", "r"), Loader=yaml.FullLoader) def log_string(str): logger.info(str) print(str) '''HYPER PARAMETER''' os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu '''CREATE DIR''' timestr = str(datetime.datetime.now().strftime('%Y-%m-%d_%H-%M')) experiment_dir = Path('./log/') experiment_dir.mkdir(exist_ok=True) experiment_dir = experiment_dir.joinpath('classification') experiment_dir.mkdir(exist_ok=True) if args.log_dir is None: experiment_dir = experiment_dir.joinpath(timestr) else: experiment_dir = experiment_dir.joinpath(args.log_dir) experiment_dir.mkdir(exist_ok=True) checkpoints_dir = experiment_dir.joinpath('checkpoints/') checkpoints_dir.mkdir(exist_ok=True) log_dir = experiment_dir.joinpath('logs/') log_dir.mkdir(exist_ok=True) '''LOG''' args = parse_args() logger = logging.getLogger("Model") logger.setLevel(logging.INFO) formatter = logging.Formatter( '%(asctime)s - %(name)s - %(levelname)s - %(message)s') file_handler = logging.FileHandler('%s/%s.txt' % (log_dir, args.model)) file_handler.setLevel(logging.INFO) file_handler.setFormatter(formatter) logger.addHandler(file_handler) log_string('PARAMETER ...') log_string(args) '''DATA LOADING''' log_string('Load dataset ...') DATA_PATH = 'data/modelnet40_normal_resampled/' TRAIN_DATASET = ModelNetDataLoader(root=DATA_PATH, npoint=args.num_point, split='train', normal_channel=args.normal) TEST_DATASET = ModelNetDataLoader(root=DATA_PATH, npoint=args.num_point, split='test', normal_channel=args.normal) trainDataLoader = torch.utils.data.DataLoader(TRAIN_DATASET, batch_size=args.batch_size, shuffle=True, num_workers=8) testDataLoader = torch.utils.data.DataLoader(TEST_DATASET, batch_size=args.batch_size, shuffle=False, num_workers=8) '''MODEL LOADING''' num_class = 40 MODEL = importlib.import_module(args.model) shutil.copy('./models/%s.py' % args.model, str(experiment_dir)) shutil.copy('./models/pointnet_util.py', str(experiment_dir)) # online network online_network = MODEL.get_model(num_class, normal_channel=args.normal).cuda() criterion = MODEL.get_loss().cuda() # predictor network device = 'cuda' if torch.cuda.is_available() else 'cpu' predictor = MLPHead( in_channels=online_network.projetion.net[-1].out_features, **config['network']['projection_head']).to(device) # target encoder target_network = MODEL.get_model(num_class, normal_channel=args.normal).cuda() # load pre-trained model if defined try: checkpoint = torch.load('checkpoints/model.pth') online_network.load_state_dict(checkpoint['online_network_state_dict']) target_network.load_state_dict(checkpoint['target_network_state_dict']) log_string('Use pretrain model') except: log_string('No existing model, starting training from scratch...') start_epoch = 0 # if args.optimizer == 'Adam': # optimizer = torch.optim.Adam( # classifier.parameters(), # lr=args.learning_rate, # betas=(0.9, 0.999), # eps=1e-08, # weight_decay=args.decay_rate # ) # else: # optimizer = torch.optim.SGD(classifier.parameters(), lr=0.01, momentum=0.9) optimizer = torch.optim.SGD( list(online_network.parameters()) + list(predictor.parameters()), **config['optimizer']['params']) trainer = BYOLTrainer(online_network=online_network, target_network=target_network, predictor=predictor, optimizer=optimizer, device=device, **config['trainer']) trainer.train_pointnet(trainDataLoader, testDataLoader)
def main(): config = yaml.load(open("./config/config.yaml", "r")) os.environ["CUDA_VISIBLE_DEVICES"] = config['CUDA_VISIBLE_DEVICES'] device = 'cuda' if torch.cuda.is_available() else 'cpu' print(f"Training with: {device}") #data_transform = get_simclr_data_transforms(**config['data_transforms']) train_source,ft_source,test_source = load_data(config) ''' train_source,_,_ = load_data(config) config_test=config config_test['dataset_path']="/home/yqliu/Dataset/CASIA-B/silhouettes" config_test['dataset']='CASIA-B' ft_source,_,test_source = load_data(config_test) print(config_test['dataset']) ''' #print(type(train_source),type(test_source)) print(train_source.__len__(),ft_source.__len__(),test_source.__len__()) train=True test=False Train_Set=Test_Set=None if train: print("Loading training data...") #Train_Set = train_source.load_all_data() if test: print("Loading test data...") #Test_Set = test_source.load_all_data() #print(type(Train_Set),type(Test_Set)) # online network #online_network = ResNet18(**config['network']).to(device) online_network = SetNet(config['network']['hidden_dim']) online_network=nn.DataParallel(online_network).cuda() pretrained_folder = config['network']['fine_tune_from'] conv_trans=MCM_NOTP(in_channels=128, out_channels=128, p=31, div=4) conv_trans=nn.DataParallel(conv_trans).cuda() #conv_trans=MCM_NOTP(in_channels=config['network']['hidden_dim'], out_channels=config['network']['hidden_dim'], p=16, div=4).to(device) TP_1=TP_FULL(hidden_dim=config['network']['hidden_dim']) TP_1=nn.DataParallel(TP_1).cuda() # load pre-trained model if defined if pretrained_folder!='None': try: checkpoints_folder = os.path.join('./runs', pretrained_folder, 'checkpoints') # load pre-trained parameters load_params = torch.load(os.path.join(os.path.join(checkpoints_folder, 'model.pth')), map_location=torch.device(torch.device(device))) online_network.load_state_dict(load_params['online_network_state_dict']) except FileNotFoundError: print("Pre-trained weights not found. Training from scratch.") # predictor network projection_1 = MLPHead(in_channels=config['network']['hidden_dim'],**config['network']['projection_head']) projection_1=nn.DataParallel(projection_1).cuda() #projection = MLPHead(in_channels=128,**config['network']['projection_head']).to(device) #predictor_1=predictor(hidden_dim=config['network']['hidden_dim']).cuda() predictor = MLPHead(in_channels=config['network']['hidden_dim'],**config['network']['projection_head']) predictor=nn.DataParallel(predictor).cuda() # target encoder #target_network = ResNet18(**config['network']).to(device) target_network = SetNet(config['network']['hidden_dim']) target_network=nn.DataParallel(target_network).cuda() TP_2=TP_FULL(hidden_dim=config['network']['hidden_dim']) TP_2=nn.DataParallel(TP_2).cuda() #predictor_2=predictor(hidden_dim=config['network']['hidden_dim']).cuda() projection_2 = MLPHead(in_channels=config['network']['hidden_dim'],**config['network']['projection_head']) projection_2=nn.DataParallel(projection_2).cuda() optimizer = torch.optim.SGD(list(online_network.parameters()) + list(conv_trans.parameters()) + list(TP_1.parameters()) + list(projection_1.parameters()) + list(predictor.parameters()),**config['optimizer']['params']) #print(type(projection_1),type(projection_2)) trainer = BYOLTrainer(online_network=online_network, target_network=target_network, conv_trans=conv_trans, TP_1=TP_1, TP_2=TP_2, projection_1=projection_1, projection_2=projection_2, predictor=predictor, optimizer=optimizer, device=device, config=config, train_source=train_source, test_source=test_source, **config['trainer']) #trainer.train_ft(train_source,maxiter=10000) #trainer.finetune(ft_source,maxiter=80000) trainer.load_model("/home/yqliu/PJs/BYOL_MCM/PyTorch-BYOL/runs/Oct20_11-37-30_123.pami.group/checkpoints/model_9000.pth") trainer.finetune(train_source,maxiter=80000)