def main(): args = parser.parse_args() assert args.n_views == 2, "Only two view training is supported. Please use --n-views 2." # check if gpu training is available if not args.disable_cuda and torch.cuda.is_available(): args.device = torch.device('cuda') cudnn.deterministic = True cudnn.benchmark = True else: args.device = torch.device('cpu') args.gpu_index = -1 dataset = ContrastiveLearningDataset(args.data) train_dataset = dataset.get_dataset(args.dataset_name, args.n_views) train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=True, drop_last=True) model = ResNetSimCLR(base_model=args.arch, out_dim=args.out_dim) optimizer = torch.optim.Adam(model.parameters(), args.lr, weight_decay=args.weight_decay) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=len(train_loader), eta_min=0, last_epoch=-1) # It’s a no-op if the 'gpu_index' argument is a negative integer or None. with torch.cuda.device(args.gpu_index): simclr = SimCLR(model=model, optimizer=optimizer, scheduler=scheduler, args=args) simclr.train(train_loader)
def main(): args = parser.parse_args() assert args.n_views == 2, "Only two view training is supported. Please use --n-views 2." # check if gpu training is available if not args.disable_cuda and torch.cuda.is_available(): args.device = torch.device('cuda') cudnn.deterministic = True cudnn.benchmark = True else: args.device = torch.device('cpu') args.gpu_index = -1 dataset = ContrastiveLearningDataset(args.data) train_dataset = dataset.get_dataset(args.dataset_name, args.n_views) train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=True, drop_last=True) model = ResNetSimCLR(base_model=args.arch, out_dim=args.out_dim).to(args.device) optimizer = torch.optim.Adam(model.parameters(), args.lr, weight_decay=args.weight_decay) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, T_max=len(train_loader), eta_min=0, last_epoch=-1) if args.resume: if os.path.isfile(args.resume): print("=> loading resumed checkpoint '{}'".format(args.resume)) checkpoint = torch.load(args.resume, map_location=torch.device('cpu')) args.start_epoch = checkpoint['epoch'] args.start_iter = checkpoint['iteration'] model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) scheduler.load_state_dict(checkpoint['scheduler']) print("=> loaded resumed checkpoint '{}' (epoch {})".format( args.resume, checkpoint['epoch'])) else: print("[Warning] no checkpoint found at '{}'".format(args.resume)) else: args.start_iter = 0 args.start_epoch = 0 # It’s a no-op if the 'gpu_index' argument is a negative integer or None. with torch.cuda.device(args.gpu_index): simclr = SimCLR(model=model, optimizer=optimizer, scheduler=scheduler, args=args) simclr.train(train_loader)
def predict(self): #Data valid_loader = self.dataset.get_data_loaders() #Model model = ResNetSimCLR(**self.config["model"]) if self.device == 'cuda': model = nn.DataParallel( model, device_ids=[i for i in range(self.config['gpu']['gpunum'])]) model = model.cuda() #print(model) model = self._load_pre_trained_weights(model) # validate the model if requested self._validate(model, valid_loader)
def __init__(self, dataset, config): self.config = config self.device = self._get_device() self.writer = SummaryWriter() self.dataset = dataset self.model = ResNetSimCLR(**self.config["model"]).to(self.device) self.criterion = self._make_loss()
def __init__(self, dataset, config): self.config = config self.device = self._get_device() self.writer = SummaryWriter() self.dataset = dataset self.batch_size = config['batch_size'] self.model = ResNetSimCLR(**self.config["model"]).to(self.device) self.loss_func = self._choose_loss()
def load_model(checkpoints_folder,device): model =ResNetSimCLR(**config['model']) model.eval() state_dict = torch.load(os.path.join(checkpoints_folder, 'model.pth'), map_location=torch.device('cpu')) model.load_state_dict(state_dict) model = model.to(device) return model
def _build_model(self): model = ResNetSimCLR(**self.config["model"]).to(self.device) model = self._load_pre_trained_weights(model) augmentor = IcmAugmentorv2(num_mech=self.config["num_mechanisms"]).to( self.device) augmentor = self._load_pre_trained_weights(augmentor, "augmentor.pth") discriminator = Discriminatorv2( num_mech=self.config["num_mechanisms"]).to(self.device) discriminator = self._load_pre_trained_weights(discriminator, "discriminator.pth") return model, augmentor, discriminator
def encode(save_root, model_file, data_folder, model_name='ca', dataset_name='celeba', batch_size=64, device='cuda:0', out_dim=256): os.makedirs(save_root, exist_ok=True) os.makedirs(data_folder, exist_ok=True) if dataset_name == 'celeba': train_loader = DataLoader(datasets.CelebA(data_folder, split='train', download=True, transform=transforms.ToTensor()), batch_size=batch_size, shuffle=False) valid_loader = DataLoader(datasets.CelebA(data_folder, split='valid', download=True, transform=transforms.ToTensor()), batch_size=batch_size, shuffle=False) elif dataset_name == 'stanfordCars': t = transforms.Compose([ transforms.Resize(512), transforms.CenterCrop(512), transforms.ToTensor(), transforms.Lambda(lambda x: x.repeat(3,1,1) if x.shape[0] == 1 else x) ]) train_data_dir = os.path.join(data_folder, 'cars_train/') train_annos = os.path.join(data_folder, 'devkit/cars_train_annos.mat') train_loader = DataLoader(CarsDataset(train_annos, train_data_dir, t), batch_size=batch_size, shuffle=False) valid_data_dir = os.path.join(data_folder, 'cars_test/') valid_annos = os.path.join(data_folder, 'devkit/cars_test_annos_withlabels.mat') valid_loader = DataLoader(CarsDataset(valid_annos, valid_data_dir, t), batch_size=batch_size, shuffle=False) elif dataset_name == 'compCars': t = transforms.Compose([ transforms.Resize(512), transforms.CenterCrop(512), transforms.ToTensor() ]) train_loader = DataLoader(CompCars(data_folder, True, t), batch_size=batch_size, shuffle=False) valid_loader = DataLoader(CompCars(data_folder, False, t), batch_size=batch_size, shuffle=False) model = ResNetSimCLR('resnet50', out_dim) model.load_state_dict(torch.load(model_file, map_location=device)) model = model.to(device) model.eval() print('Starting on training data') train_encodings = [] for x, _ in train_loader: x = x.to(device) h, _ = model(x) train_encodings.append(h.cpu().detach()) torch.save(torch.cat(train_encodings, dim=0), os.path.join(save_root, f'{dataset_name}-{model_name}model-train_encodings.pt')) print('Starting on validation data') valid_encodings = [] for x, _ in valid_loader: x = x.to(device) h, _ = model(x) if len(h.shape) == 1: h = h.unsqueeze(0) valid_encodings.append(h.cpu().detach()) torch.save(torch.cat(valid_encodings, dim=0), os.path.join(save_root, f'{dataset_name}-{model_name}model-valid_encodings.pt'))
def _save_dt_features(self): from eval.feature_saver import LvisSaver import torch config = self.config if self.config['hyperbolic']: from models.hyperbolic_resnet import HResNetSimCLR model = HResNetSimCLR(config['model']['base_model'], config['model']['out_dim']) else: from models.resnet_simclr import ResNetSimCLR model = ResNetSimCLR(config['model']['base_model'], config['model']['out_dim']) state_dict = torch.load(os.path.join(self.run_path, self.model_ckpt)) # , map_location=device)i model.load_state_dict(state_dict) model.eval() saver = LvisSaver(model, self.lvis_dt, GT_FEATS) saver.save()
def main(): config_path = "config.json" args = process_config(config_path) # check if gpu training is available if not args.disable_cuda and torch.cuda.is_available(): args.device = torch.device('cuda') cudnn.deterministic = True cudnn.benchmark = True else: args.device = torch.device('cpu') args.gpu_index = -1 log_dir = time.ctime() os.makedirs(log_dir) model = ResNetSimCLR(base_model=args.arch, out_dim=args.out_dim) load_checkpoint(model, args.simclr.pretrained_path) # run_simclr(args, model, log_dir) run_la(args, model, log_dir)
def _init_model_and_optimizer(self): if self.config['hyperbolic']: model = HResNetSimCLR(**self.config["model"]).to(self.device) else: model = ResNetSimCLR(**self.config["model"]).to(self.device) model, loaded_iter = self._load_pre_trained_weights(model) # if self.config['hyperbolic']: # optimizer = geoopt.optim.RiemannianAdam( # [p for p in self.model.parameters() if p.requires_grad], # 1e-4, weight_decay=eval(self.config['weight_decay'])) # else: lr = float(self.config['lr']) optimizer = torch.optim.Adam( [p for p in model.parameters() if p.requires_grad], lr, weight_decay=eval(self.config['weight_decay'])) self.model = model self.optimizer = optimizer return loaded_iter
def main(): args = parser.parse_args() assert args.n_views == 2, "Only two view training is supported. Please use --n-views 2." # check if gpu training is available if not args.disable_cuda and torch.cuda.is_available(): args.device = torch.device('cuda') cudnn.deterministic = True cudnn.benchmark = True else: args.device = torch.device('cpu') args.gpu_index = -1 set_random_seed(args.seed) train_loader, valid_loader = get_dataloaders(args) if args.mode == 'simclr': model = ResNetSimCLR(base_model=args.arch, out_dim=args.out_dim) trainer_class = SimCLRTrainer elif args.mode == 'supervised': model = ResNetSimCLR(base_model=args.arch, out_dim=len(train_loader.dataset.classes)) trainer_class = SupervisedTrainer else: raise InvalidTrainingMode() if args.optimizer_mode == 'simclr': optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=len(train_loader), eta_min=0, last_epoch=-1) else: optimizer = torch.optim.SGD(model.parameters(), lr=args.lr, momentum=0.9, weight_decay=args.weight_decay) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=200) # It’s a no-op if the 'gpu_index' argument is a negative integer or None. with torch.cuda.device(args.gpu_index): trainer = trainer_class(model=model, optimizer=optimizer, scheduler=scheduler, args=args) trainer.train(train_loader, valid_loader)
def run(config): model = ResNetSimCLR('resnet50', config.out_dim) model.load_state_dict( torch.load(config.model_file, map_location=config.device)) model = model.to(config.device) clf = nn.Linear(2048, 196) full = FineTuner(model, clf) optim = torch.optim.Adam(list(model.parameters()) + list(clf.parameters())) objective = nn.CrossEntropyLoss() t = T.Compose([ T.Resize(512), T.CenterCrop(512), T.ToTensor(), T.Lambda(lambda x: x.repeat(3, 1, 1) if x.shape[0] == 1 else x) ]) train_data_dir = os.path.join(config.data_root, 'cars_train/') train_annos = os.path.join(config.data_root, 'devkit/cars_train_annos.mat') train_dataset = StanfordCarsMini(train_annos, train_data_dir, t) train_loader = DataLoader(train_dataset, batch_size=config.batch_size, shuffle=True) valid_data_dir = os.path.join(config.data_root, 'cars_test/') valid_annos = os.path.join(config.data_root, 'devkit/cars_test_annos_withlabels.mat') valid_dataset = CarsDataset(valid_annos, valid_data_dir, t) valid_loader = DataLoader(valid_dataset, batch_size=config.batch_size) solver = CESolver(full, train_loader, valid_loader, config.save_root, name=config.name, device=config.device) solver.train(config.num_epochs)
def eval_frozen(self): train_loader, val_loader, num_classes = self.dataset.get_dataset_eval() model = ResNetSimCLR(**self.config["model"]).to(self.device) model = self._load_pre_trained_weights(model) model.to(self.device) model.eval() lineal_classifier = nn.Linear(model.l1.in_features, num_classes) lineal_classifier.to(self.device) optimizer = torch.optim.SGD(lineal_classifier.parameters(), 1e-3, weight_decay=eval( self.config['weight_decay'])) epochs = self.config['epochs'] scheduler = torch.optim.lr_scheduler.MultiStepLR( optimizer, milestones=[int(0.5 * epochs), int(0.8 * epochs)], gamma=0.1, last_epoch=-1) criterion = nn.CrossEntropyLoss() best_acc = 0 for epoch in range(epochs): print('=' * 20) print(f'Epoch: {epoch + 1} / {epochs}') top1 = 0 running_loss = 0 n = 0 lineal_classifier.train() for idx, (img, lab) in enumerate(train_loader): B = img.size(0) img = img.to(self.device) lab = lab.to(self.device) optimizer.zero_grad() loss, top1_batch = self._step_eval_train( model, lineal_classifier, img, lab) loss.backward() optimizer.step() top1 += (x.argmax(dim=1) == lab).sum().item() running_loss += loss.item() * B n += B print('Training {}/{} - Loss: {:.2f} - top 1: {:.2f}'.format( idx + 1, len(train_loader), running_loss / n, 100 * top1 / n), end='\r') print('\n') top1 = 0 running_loss = 0 n = 0 lineal_classifier.eval() for idx, (img, lab) in enumerate(val_loader): B = img.size(0) img = img.to(self.device) lab = lab.to(self.device) loss, top1_batch = self._step_eval_eval( model, lineal_classifier, img, lab) top1 += top1_batch running_loss += loss.item() * B n += B print('Val {}/{} - Loss: {:.2f} - top 1: {:.2f}'.format( idx + 1, len(val_loader), running_loss / n, 100 * top1 / n), end='\r') print('\n') if best_acc < top1: best_acc = top1 print(f'Best ACC: {best_acc * 100 / n}')
def train(self): train_loader, valid_loader = self.dataset.get_data_loaders() model = ResNetSimCLR(**self.config["model"]).to(self.device) model = self._load_pre_trained_weights(model) optimizer = torch.optim.Adam(model.parameters(), 3e-4, weight_decay=eval( self.config['weight_decay'])) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, T_max=self.config['epochs'], eta_min=0, last_epoch=-1) if apex_support and self.config['fp16_precision']: model, optimizer = amp.initialize(model, optimizer, opt_level='O2', keep_batchnorm_fp32=True) model_checkpoints_folder = os.path.join(self.writer.log_dir, 'checkpoints') # save config file _save_config_file(model_checkpoints_folder) n_iter = 0 valid_n_iter = 0 best_valid_loss = np.inf for epoch_counter in range(self.config['epochs']): for (xis, xjs), _ in train_loader: optimizer.zero_grad() xis = xis.to(self.device) xjs = xjs.to(self.device) loss = self._step(model, xis, xjs, n_iter) if n_iter % self.config['log_every_n_steps'] == 0: self.writer.add_scalar('train_loss', loss, global_step=n_iter) if apex_support and self.config['fp16_precision']: with amp.scale_loss(loss, optimizer) as scaled_loss: scaled_loss.backward() else: loss.backward() optimizer.step() n_iter += 1 # validate the model if requested if epoch_counter % self.config['eval_every_n_epochs'] == 0: valid_loss = self._validate(model, valid_loader) if valid_loss < best_valid_loss: # save the model weights best_valid_loss = valid_loss torch.save( model.state_dict(), os.path.join(model_checkpoints_folder, 'model.pth')) self.writer.add_scalar('validation_loss', valid_loss, global_step=valid_n_iter) valid_n_iter += 1 # warmup for the first 10 epochs if epoch_counter >= 10: scheduler.step() self.writer.add_scalar('cosine_lr_decay', scheduler.get_lr()[0], global_step=n_iter)
def run(config): model = ResNetSimCLR('resnet50', config.out_dim) model.load_state_dict( torch.load(config.model_file, map_location=config.device)) model = model.to(config.device) clf = nn.Linear(2048, 196) train_data_dir = os.path.join(config.data_root, 'cars_train/') train_annos = os.path.join(config.data_root, 'devkit/cars_train_annos.mat') valid_data_dir = os.path.join(config.data_root, 'cars_test/') valid_annos = os.path.join(config.data_root, 'devkit/cars_test_annos_withlabels.mat') if config.encodings_file_prefix: train_dataset = EncodedStanfordCarsDataset( train_annos, config.encodings_file_prefix + '-train_encodings.pt') train_loader = DataLoader(train_dataset, batch_size=config.encodings_batch_size, shuffle=True) valid_dataset = EncodedStanfordCarsDataset( valid_annos, config.encodings_file_prefix + '-valid_encodings.pt') valid_loader = DataLoader(valid_dataset, batch_size=config.encodings_batch_size) tmp_clf = nn.Linear(2048, 196) clf_solver = CESolver(tmp_clf, train_loader, valid_loader, config.save_root, name=config.name + '-clf', device=config.device) clf_solver.train(config.encodings_num_epochs) clf_filename = os.path.join(config.save_root, f'{config.name}-clf.pth') clf.load_state_dict( torch.load(clf_filename, map_location=config.device)) full = FineTuner(model, clf) t = T.Compose([ T.Resize(512), T.CenterCrop(512), T.ToTensor(), T.Lambda(lambda x: x.repeat(3, 1, 1) if x.shape[0] == 1 else x) ]) train_dataset = StanfordCarsMini(train_annos, train_data_dir, t) train_loader = DataLoader(train_dataset, batch_size=config.batch_size, shuffle=True) valid_dataset = CarsDataset(valid_annos, valid_data_dir, t) valid_loader = DataLoader(valid_dataset, batch_size=config.batch_size) solver = CESolver(full, train_loader, valid_loader, config.save_root, name=config.name, device=config.device) # full = full.to(config.device) # print(solver.validate(full, nn.CrossEntropyLoss())) solver.train(config.num_epochs)
def main(): args = parser.parse_args() assert args.n_views == 2, "Only two view training is supported. Please use --n-views 2." # check if gpu training is available if not args.disable_cuda and torch.cuda.is_available(): args.device = torch.device('cuda') cudnn.deterministic = True cudnn.benchmark = True else: args.device = torch.device('cpu') args.gpu_index = -1 if args.mode == 'simclr': dataset = ContrastiveLearningDataset(args.data) train_dataset = dataset.get_dataset(args.dataset_name, args.n_views, train=True) model = ResNetSimCLR(base_model=args.arch, out_dim=args.out_dim) trainer_class = SimCLRTrainer elif args.mode == 'supervised': dataset = SupervisedLearningDataset(args.data) train_dataset = dataset.get_dataset(args.dataset_name, args.supervised_augments, train=True) model = ResNetSimCLR(base_model=args.arch, out_dim=len(train_dataset.classes)) trainer_class = SupervisedTrainer else: raise InvalidTrainingMode() if args.target_shuffle is not None: random.seed(args.target_shuffle) random.shuffle(train_dataset.targets) checkpoints = [] for root, dirs, files in os.walk( os.path.join('experiments', args.experiment_group, 'wandb')): for file in files: if file == args.estimate_checkpoint: checkpoints += [os.path.join(root, file)] set_random_seed(args.seed) sample_indices = torch.randint(len(train_dataset), size=(args.batch_size * args.estimate_batches, )) # It’s a no-op if the 'gpu_index' argument is a negative integer or None. estimated_prob, estimated_argmax = [], [] with torch.cuda.device(args.gpu_index): for file in checkpoints: state = torch.load(file) model.load_state_dict(state['model']) model.eval() trainer = trainer_class(model=model, optimizer=None, scheduler=None, args=args) checkpoint_prob, checkpoint_argmax = [], [] for i in range(args.estimate_batches): if args.fixed_augments: set_random_seed(args.seed) if args.mode == 'simclr': images = [[], []] for index in sample_indices[i:i + args.batch_size]: example = train_dataset[index][0] images[0] += [example[0]] images[1] += [example[1]] images[0] = torch.stack(images[0], dim=0) images[1] = torch.stack(images[1], dim=0) labels = None elif args.mode == 'supervised': images, labels = [], [] for index in sample_indices[i:i + args.batch_size]: example = train_dataset[index] images += [example[0]] labels += [example[1]] images = torch.stack(images, dim=0) labels = torch.tensor(labels, dtype=torch.long) with torch.no_grad(): logits, labels = trainer.calculate_logits(images, labels) prob = torch.softmax(logits, dim=1)[torch.arange(labels.shape[0]), labels] checkpoint_prob += [prob.detach().cpu()] argmax = (torch.argmax(logits, dim=1) == labels).to(torch.int) checkpoint_argmax += [argmax.detach().cpu()] checkpoint_prob = torch.cat(checkpoint_prob, dim=0) estimated_prob += [checkpoint_prob] checkpoint_argmax = torch.cat(checkpoint_argmax, dim=0) estimated_argmax += [checkpoint_argmax] estimated_prob = torch.stack(estimated_prob, dim=0) estimated_argmax = torch.stack(estimated_argmax, dim=0) torch.save( { 'indices': sample_indices, 'prob': estimated_prob, 'argmax': estimated_argmax }, os.path.join('experiments', args.experiment_group, args.out_file))
test_loss = (test_loss * i + criterion(out, labels).item()) / (i + 1) res = torch.argmax(out, 1) correct += (res == labels).sum().item() tot += labels.size(0) return 100 - 100 * correct / tot, test_loss config_path = "./best_run/config.yaml" model_path = "./best_run/model.pth" data_path = "./data" config = yaml.load(open(config_path, "r"), Loader=yaml.FullLoader) device = 'cuda:0' if torch.cuda.is_available() else 'cpu' model = ResNetSimCLR(**config["model"]).to(device) model.load_state_dict(torch.load(model_path)) lr = 0.001 num_epoch = 90 batch_size = 512 num_classes = 10 weight_decay = 1e-6 transform = transforms.Compose([ transforms.ToTensor(), ]) linear = Linear_Model(512, num_classes) # Linear_Model or MLP linear.to(device)
def train(self): #Data train_loader, valid_loader = self.dataset.get_data_loaders() #Model model = ResNetSimCLR(**self.config["model"]) if self.device == 'cuda': model = nn.DataParallel(model, device_ids=[i for i in range(self.config['gpu']['gpunum'])]) #model = model.to(self.device) model = model.cuda() print(model) model = self._load_pre_trained_weights(model) each_epoch_steps = len(train_loader) total_steps = each_epoch_steps * self.config['train']['epochs'] warmup_steps = each_epoch_steps * self.config['train']['warmup_epochs'] scaled_lr = eval(self.config['train']['lr']) * self.batch_size / 256. optimizer = torch.optim.Adam( model.parameters(), scaled_lr, weight_decay=eval(self.config['train']['weight_decay'])) ''' optimizer = LARS(params=model.parameters(), lr=eval(self.config['train']['lr']), momentum=self.config['train']['momentum'], weight_decay=eval(self.config['train']['weight_decay'], eta=0.001, max_epoch=self.config['train']['epochs']) ''' # scheduler during warmup stage lambda1 = lambda epoch:epoch*1.0 / int(warmup_steps) scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda1) if apex_support and self.config['train']['fp16_precision']: model, optimizer = amp.initialize(model, optimizer, opt_level='O2', keep_batchnorm_fp32=True) model_checkpoints_folder = os.path.join(self.writer.log_dir, 'checkpoints') # save config file _save_config_file(model_checkpoints_folder) n_iter = 0 valid_n_iter = 0 best_valid_loss = np.inf lr = eval(self.config['train']['lr']) end = time.time() batch_time = AverageMeter() data_time = AverageMeter() losses = AverageMeter() for epoch_counter in range(self.config['train']['epochs']): model.train() for i, ((xis, xjs), _) in enumerate(train_loader): data_time.update(time.time() - end) optimizer.zero_grad() xis = xis.cuda() xjs = xjs.cuda() loss = self._step(model, xis, xjs, n_iter) #print("Loss: ",loss.data.cpu()) losses.update(loss.item(), 2 * xis.size(0)) # measure elapsed time batch_time.update(time.time() - end) end = time.time() print('Epoch: [{epoch}][{step}/{each_epoch_steps}] Loss {loss.val:.4f} Avg Loss {loss.avg:.4f} DataTime {datatime.val:.4f} BatchTime {batchtime.val:.4f} LR {lr})'.format(epoch=epoch_counter, step=i, each_epoch_steps=each_epoch_steps, loss=losses, datatime=data_time, batchtime=batch_time, lr=lr)) if n_iter % self.config['train']['log_every_n_steps'] == 0: self.writer.add_scalar('train_loss', loss, global_step=n_iter) if apex_support and self.config['train']['fp16_precision']: with amp.scale_loss(loss, optimizer) as scaled_loss: scaled_loss.backward() else: loss.backward() optimizer.step() n_iter += 1 #adjust lr if n_iter == warmup_steps: # scheduler after warmup stage scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=total_steps-warmup_steps, eta_min=0, last_epoch=-1) scheduler.step() lr = scheduler.get_lr()[0] self.writer.add_scalar('cosine_lr_decay', scheduler.get_lr()[0], global_step=n_iter) sys.stdout.flush() # validate the model if requested if epoch_counter % self.config['train']['eval_every_n_epochs'] == 0: valid_loss = self._validate(model, valid_loader) if valid_loss < best_valid_loss: # save the model weights best_valid_loss = valid_loss torch.save(model.state_dict(), os.path.join(model_checkpoints_folder, 'model.pth')) self.writer.add_scalar('validation_loss', valid_loss, global_step=valid_n_iter) valid_n_iter += 1
# dataset_train, batch_size=args.batch_size, shuffle=True, # num_workers=args.workers, pin_memory=True, drop_last=True) # build model if args.model == 'cnn' and args.dataset == 'cifar': net_glob = CNNCifar(args=args).to(args.device) elif args.model == 'cnn' and args.dataset == 'mnist': net_glob = CNNMnist(args=args).to(args.device) elif args.model == 'mlp': len_in = 1 for x in img_size: len_in *= x net_glob = MLP(dim_in=len_in, dim_hidden=200, dim_out=args.num_classes).to(args.device) elif args.model == 'resnet': net_glob = ResNetSimCLR(base_model=args.arch, out_dim=args.out_dim).to(args.device) else: exit('Error: unrecognized model') net_glob.train() # copy weights w_glob = net_glob.state_dict() dict_glob = {} # training loss_train = [] cv_loss, cv_acc = [], [] val_loss_pre, counter = 0, 0 net_best = None best_loss = None
def main_worker(gpu, ngpus_per_node, args): args.gpu = gpu # suppress printing if not master if args.multiprocessing_distributed and args.gpu != 0: def print_pass(*args): pass builtins.print = print_pass if args.gpu is not None: print(args.gpu) print("Use GPU: {} for training".format(args.gpu)) if args.distributed: if args.dist_url == "env://" and args.rank == -1: args.rank = int(os.environ["RANK"]) if args.multiprocessing_distributed: # For multiprocessing distributed training, rank needs to be the # global rank among all the processes args.rank = args.rank * ngpus_per_node + gpu dist.init_process_group(backend=args.dist_backend, init_method=args.dist_url, world_size=args.world_size, rank=args.rank) # create model print("=> creating model '{}'".format(args.arch)) model = ResNetSimCLR(base_model=args.arch, out_dim=args.out_dim).to(args.gpu) if args.distributed: # For multiprocessing distributed, DistributedDataParallel constructor # should always set the single device scope, otherwise, # DistributedDataParallel will use all available devices. if args.gpu is not None: torch.cuda.set_device(args.gpu) model.cuda(args.gpu) # When using a single GPU per process and per # DistributedDataParallel, we need to divide the batch size # ourselves based on the total number of GPUs we have args.batch_size = int(args.batch_size / ngpus_per_node) args.num_workers = int( (args.num_workers + ngpus_per_node - 1) / ngpus_per_node) model = nn.SyncBatchNorm.convert_sync_batchnorm(model) model = torch.nn.parallel.DistributedDataParallel( model, device_ids=[args.gpu]) else: model.cuda() # DistributedDataParallel will divide and allocate batch_size to all # available GPUs if device_ids are not set model = torch.nn.parallel.DistributedDataParallel(model) elif args.gpu is not None: torch.cuda.set_device(args.gpu) model = model.cuda(args.gpu) # comment out the following line for debugging #raise NotImplementedError("Only DistributedDataParallel is supported.") #else: # AllGather implementation (batch shuffle, queue update, etc.) in # this code only supports DistributedDataParallel. #raise NotImplementedError("Only DistributedDataParallel is supported.") # Data loader train_loader, train_sampler = data_loader(args.dataset, args.data_path, args.batch_size, args.num_workers, download=args.download, distributed=args.distributed, supervised=False) #optimizer = torch.optim.Adam(model.parameters(), 3e-4, weight_decay=args.weight_decay) optimizer = torch.optim.SGD(model.parameters(), lr=args.lr, momentum=0.9, weight_decay=args.weight_decay) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=args.epochs, eta_min=0, last_epoch=-1) criterion = NTXentLoss(args.gpu, args.batch_size, args.temperature, True).cuda(args.gpu) if args.resume: if os.path.isfile(args.resume): print("=> loading checkpoint '{}'".format(args.resume)) if args.gpu is None: checkpoint = torch.load(args.resume) else: loc = 'cuda:{}'.format(args.gpu) checkpoint = torch.load(args.resume, map_location=loc) args.start_epoch = checkpoint['epoch'] model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) print("=> loaded checkpoint '{}' (epoch {})".format( args.resume, checkpoint['epoch'])) else: print("=> no checkpoint found at '{}'".format(args.resume)) if apex_support and args.fp16_precision: model, optimizer = amp.initialize(model, optimizer, opt_level='O2', keep_batchnorm_fp32=True) cudnn.benchmark = True train(model, train_loader, train_sampler, criterion, optimizer, scheduler, args, ngpus_per_node)
train_dataset = datasets.STL10('./data', split='train+unlabeled', download=True, transform=DataTransform(data_augment)) # train_dataset = datasets.Caltech101(root='./data', target_type="category", transform=transforms.ToTensor(), # target_transform=None, download=True) train_loader = DataLoader(train_dataset, batch_size=batch_size, num_workers=config['num_workers'], drop_last=True, shuffle=True) # model = Encoder(out_dim=out_dim) model = ResNetSimCLR(base_model=config["base_convnet"], out_dim=out_dim) train_gpu = torch.cuda.is_available() print("Is gpu available:", train_gpu) # moves the model parameters to gpu if train_gpu: model.cuda() criterion = torch.nn.CrossEntropyLoss(reduction='sum') optimizer = optim.Adam(model.parameters(), 3e-4) train_writer = SummaryWriter() sim_func_dim1, sim_func_dim2 = get_similarity_function(use_cosine_similarity)
def train(self): train_loader, valid_loader = self.dataset.get_data_loaders() model = ResNetSimCLR(**self.config["model"]).to(self.device) #just a resnet backbone model = self._load_pre_trained_weights(model) #checkpoints (shall we convert TF checkpoint in to Torch and train or train from the scratch since we have f*****g lot?) optimizer = torch.optim.Adam(model.parameters(), 3e-4, weight_decay=eval(self.config['weight_decay'])) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=len(train_loader), eta_min=0, last_epoch=-1) #learning rate shedulers (let's use as it is) if apex_support and self.config['fp16_precision']: model, optimizer = amp.initialize(model, optimizer, opt_level='O2', keep_batchnorm_fp32=True) model_checkpoints_folder = os.path.join(self.writer.log_dir, 'checkpoints') # save config file _save_config_file(model_checkpoints_folder) n_iter = 0 valid_n_iter = 0 best_valid_loss = np.inf for epoch_counter in range(self.config['epochs']): #start training for (x, y) in train_loader: #dataset optimizer.zero_grad() x = x.to(self.device) #in SimCLR we calculate the loss with two augmentation versions y = y.to(self.device) loss = self._step(model, x, y) if n_iter % self.config['log_every_n_steps'] == 0: self.writer.add_scalar('train_loss', loss, global_step=n_iter) if apex_support and self.config['fp16_precision']: with amp.scale_loss(loss, optimizer) as scaled_loss: scaled_loss.backward() else: loss.backward() optimizer.step() n_iter += 1 # validate the model if requested if epoch_counter % self.config['eval_every_n_epochs'] == 0: valid_loss = self._validate(model, valid_loader) print('Epoch:',epoch_counter,' ---',' validation_loss:',valid_loss) if valid_loss < best_valid_loss: # save the model weights best_valid_loss = valid_loss torch.save(model.state_dict(), os.path.join(model_checkpoints_folder, 'model.pth')) self.writer.add_scalar('validation_loss', valid_loss, global_step=valid_n_iter) valid_n_iter += 1 # warmup for the first 10 epochs if epoch_counter >= 10: scheduler.step() self.writer.add_scalar('cosine_lr_decay', scheduler.get_lr()[0], global_step=n_iter)
def train(self): train_loader, valid_loader = self.dataset.get_data_loaders() model = ResNetSimCLR(**self.config["model"]).to(self.device) model = self._load_pre_trained_weights(model) criterion = nn.CrossEntropyLoss() # loss function optimizer = torch.optim.Adam(model.parameters(), 3e-4, weight_decay=eval( self.config['weight_decay'])) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, T_max=self.config['epochs'], eta_min=0, last_epoch=-1) if apex_support and self.config['fp16_precision']: model, optimizer = amp.initialize(model, optimizer, opt_level='O2', keep_batchnorm_fp32=True) model_checkpoints_folder = os.path.join( '/home/zhangchunhui/MedicalAI/USCL/checkpoints_multi_aug', 'checkpoint_' + str(self.Checkpoint_Num)) # save config file _save_config_file(model_checkpoints_folder) start_time = time.time() end_time = time.time() valid_n_iter = 0 best_valid_loss = np.inf for epoch in range(self.config['epochs']): for i, data in enumerate(train_loader, 1): # forward # mixupimg1, label1, mixupimg2, label2, original img1, original img2 xis, labelis, xjs, labeljs, imgis, imgjs = data # N samples of left branch, N samples of right branch xis = xis.to(self.device) xjs = xjs.to(self.device) ####### 1-Semi-supervised hi, xi, outputis = model(xis) hj, xj, outputjs = model(xjs) labelindexi, labelindexj = FindNotX( labelis.tolist(), 9999), FindNotX(labeljs.tolist(), 9999) # X=9999=no label lossi = criterion(outputis[labelindexi], labelis.to(self.device)[labelindexi]) lossj = criterion(outputjs[labelindexj], labeljs.to(self.device)[labelindexj]) # lumbda1=lumbda2 # small value is better lumbda1, lumbda2 = self.lumbda1, self.lumbda2 # small value is better loss = self._step(model, xis, xjs) + lumbda1 * lossi + lumbda2 * lossj ######################################################################################################## ####### 2-Self-supervised # loss = self._step(model, xis, xjs) ######################################################################################################## # backward optimizer.zero_grad() if apex_support and self.config['fp16_precision']: with amp.scale_loss(loss, optimizer) as scaled_loss: scaled_loss.backward() else: loss.backward() # update weights optimizer.step() if i % self.config['log_every_n_steps'] == 0: # self.writer.add_scalar('train_loss', loss, global_step=i) start_time, end_time = end_time, time.time() print( "\nTraining:Epoch[{:0>3}/{:0>3}] Iteration[{:0>3}/{:0>3}] Loss: {:.4f} Time: {:.2f}s" .format(epoch + 1, self.config['epochs'], i, len(train_loader), loss, end_time - start_time)) # validate the model if requested if epoch % self.config['eval_every_n_epochs'] == 0: start_time = time.time() valid_loss = self._validate(model, valid_loader) end_time = time.time() if valid_loss < best_valid_loss: # save the model weights best_valid_loss = valid_loss torch.save( model.state_dict(), os.path.join(model_checkpoints_folder, 'best_model.pth')) print( "Valid:\t Epoch[{:0>3}/{:0>3}] Iteration[{:0>3}/{:0>3}] Loss: {:.4f} Time: {:.2f}s" .format(epoch + 1, self.config['epochs'], len(valid_loader), len(valid_loader), valid_loss, end_time - start_time)) # self.writer.add_scalar('validation_loss', valid_loss, global_step=valid_n_iter) valid_n_iter += 1 print('Learning rate this epoch:', scheduler.get_last_lr()[0]) # python >=3.7 # print('Learning rate this epoch:', scheduler.base_lrs[0]) # python 3.6 # warmup for the first 10 epochs if epoch >= 10: scheduler.step()
out_dim = config['out_dim'] temperature = config['temperature'] use_cosine_similarity = config['use_cosine_similarity'] data_augment = get_simclr_data_transform(s=config['s'], crop_size=96) train_dataset = datasets.STL10('./data', split='train+unlabeled', download=True, transform=DataTransform(data_augment)) train_loader, valid_loader = get_train_validation_data_loaders( train_dataset, **config) # model = Encoder(out_dim=out_dim) model = ResNetSimCLR(base_model=config["base_convnet"], out_dim=out_dim) if config['continue_training']: checkpoints_folder = os.path.join('./runs', config['continue_training'], 'checkpoints') state_dict = torch.load(os.path.join(checkpoints_folder, 'model.pth')) model.load_state_dict(state_dict) print("Loaded pre-trained model with success.") train_gpu = torch.cuda.is_available() print("Is gpu available:", train_gpu) # moves the model parameters to gpu if train_gpu: model = model.cuda()
def main(): args = parser.parse_args() assert args.n_views == 2, "Only two view training is supported. Please use --n-views 2." # check if gpu training is available if not args.disable_cuda and torch.cuda.is_available(): args.device = torch.device('cuda') cudnn.deterministic = True cudnn.benchmark = True torch.multiprocessing.set_start_method('spawn') else: args.device = torch.device('cpu') args.gpu_index = -1 if args.eval: # Load pretrained model and cifar10 cifar_transforms = transforms.Compose( [transforms.Resize(96), transforms.ToTensor()]) dataset = CIFAR10(root='datasets/cifar10', download=True, transform=cifar_transforms) train_dataset = CIFAR10(root='datasets/cifar10', train=True, transform=cifar_transforms) valid_dataset = CIFAR10(root='datasets/cifar10', train=False, transform=cifar_transforms) train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.workers) valid_loader = torch.utils.data.DataLoader(valid_dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.workers) if not args.baseline: model = ResNetBertSimCLR(base_model=args.arch, out_dim=args.out_dim) checkpoint = torch.load(args.saved_path) model.load_state_dict(checkpoint['state_dict']) else: model = ResNetSimCLR(base_model=args.arch, out_dim=args.out_dim) checkpoint = torch.load(args.saved_path) model.load_state_dict(checkpoint['state_dict']) classifier_model = torch.nn.Sequential(torch.nn.Linear(128, 10)) classifier_optimizer = torch.optim.Adam(classifier_model.parameters(), args.lr, weight_decay=args.weight_decay) optimizer = None scheduler = None else: # Load BertSimCLR and coco dataset dataset = ContrastiveLearningDataset(args.data) train_dataset = dataset.get_dataset(args.dataset_name, args.n_views) valid_dataset = dataset.get_dataset(args.dataset_name + 'valid', args.n_views) if not args.baseline: train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=True, drop_last=True, collate_fn=coco_collate_fn) valid_loader = torch.utils.data.DataLoader( valid_dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=True, drop_last=True, collate_fn=coco_collate_fn) model = ResNetBertSimCLR(base_model=args.arch, out_dim=args.out_dim) else: train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=True, drop_last=True) valid_loader = torch.utils.data.DataLoader( valid_dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=True, drop_last=True) model = ResNetSimCLR(base_model=args.arch, out_dim=args.out_dim) optimizer = torch.optim.Adam(model.parameters(), args.lr, weight_decay=args.weight_decay) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, T_max=len(train_loader), eta_min=0, last_epoch=-1) # emergency fix """ model.to(args.device) checkpoint = torch.load('/home/gabriel/Desktop/Projects/SimCLR/runs/nextbase/checkpoint_0018.pth.tar', map_location="cuda:0") model_state = checkpoint['state_dict']#.to(args.device) opt_state = checkpoint['optimizer']#.to(args.device) model.load_state_dict(model_state) optimizer.load_state_dict(opt_state) model.to(args.device) """ classifier_model = None classifier_optimizer = None #print(model.visual_backbone) data_loaders = {"train": train_loader, "val": valid_loader} # It’s a no-op if the 'gpu_index' argument is a negative integer or None. with torch.cuda.device(args.gpu_index): if not args.baseline: simclr = BertSimCLR(model=model, optimizer=optimizer, scheduler=scheduler, classifier_model=classifier_model, classifier_optimizer=classifier_optimizer, args=args) if args.eval: simclr.train_linear_classifier(args.epochs, data_loaders) else: simclr.train(data_loaders) else: simclr = SimCLR(model=model, optimizer=optimizer, scheduler=scheduler, classifier_model=classifier_model, classifier_optimizer=classifier_optimizer, args=args) if args.eval: simclr.train_linear_classifier(args.epochs, data_loaders) else: simclr.train(data_loaders)
def evaluation(checkpoints_folder, config, device): model = ResNetSimCLR(**config['model']) model.eval() model.load_state_dict( torch.load(os.path.join(checkpoints_folder, 'model.pth'))) model = model.to(device) train_set = torchvision.datasets.CIFAR10( root='../data/CIFAR10', transform=transforms.Compose([ transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]), train=True, download=True) test_set = torchvision.datasets.CIFAR10(root='../data/CIFAR10', transform=transforms.Compose([ transforms.ToTensor(), transforms.Normalize( mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]), train=False, download=True) # num_train = len(train_dataset) # indices = list(range(num_train)) # np.random.shuffle(indices) # # split = int(np.floor(0.05 * num_train)) # train_idx, test_idx = indices[split:], indices[:split] # # # define samplers for obtaining training and validation batches # train_sampler = torch.utils.data.sampler.SubsetRandomSampler(train_idx) # test_sampler = torch.utils.data.sampler.SubsetRandomSampler(test_idx) # ?????sampler???????????? # train_loader = torch.utils.data.DataLoader(train_set, batch_size=config['batch_size'], drop_last=True, shuffle=True) # # test_loader = torch.utils.data.DataLoader(test_set, batch_size=config['batch_size'], drop_last=True, shuffle=True) train_loader = torch.utils.data.DataLoader(train_set, batch_size=48, drop_last=True, shuffle=True) test_loader = torch.utils.data.DataLoader(test_set, batch_size=48, drop_last=True, shuffle=True) X_train_feature = [] label_train = [] for data in train_loader: x, y = data x = x.to(device) features, _ = model(x) X_train_feature.extend(features.cpu().detach().numpy()) label_train.extend(y.cpu().detach().numpy()) X_train_feature = np.array(X_train_feature) label_train = np.array(label_train) X_test_feature = [] label_test = [] for data in test_loader: x, y = data x = x.to(device) features, _ = model(x) X_test_feature.extend(features.cpu().detach().numpy()) label_test.extend(y.cpu().detach().numpy()) X_test_feature = np.array(X_test_feature) label_test = np.array(label_test) scaler = preprocessing.StandardScaler() print('ok') scaler.fit(X_train_feature) # print(X_test_feature.shape) # print(y_test.shape) linear_model_eval(scaler.transform(X_train_feature), label_train, scaler.transform(X_test_feature), label_test)
def train(self): train_loader, valid_loader = self.dataset.get_data_loaders() model = ResNetSimCLR(**self.config["model"]).to(self.device) model = self._load_pre_trained_weights(model) if self.augmentor_type == "cnn": if self.config["normalization_type"] == "original": augmentor = LpAugmentor( clip=self.config["augmentor_clip_output"]) augmentor.to(self.device) elif self.config["normalization_type"] == "spectral": augmentor = LpAugmentorSpecNorm( clip=self.config["augmentor_clip_output"]) augmentor.to(self.device) else: raise ValueError("Unregonized normalization type: {}".format( self.config["normalization_type"])) elif self.augmentor_type == "style_transfer": augmentor = LpAugmentorStyleTransfer( clip=self.config["augmentor_clip_output"]) augmentor.to(self.device) elif self.augmentor_type == "transformer": augmentor = LpAugmentorTransformer( clip=self.config["augmentor_clip_output"]) augmentor.to(self.device) else: raise ValueError("Unrecognized augmentor type: {}".format( self.augmentor_type)) augmentor_optimizer = torch.optim.Adam(augmentor.parameters(), 3e-4) augmentor_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( augmentor_optimizer, T_max=len(train_loader), eta_min=0, last_epoch=-1) optimizer = torch.optim.Adam( list(model.parameters()), 3e-4, weight_decay=eval(self.config["weight_decay"]), ) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, T_max=len(train_loader), eta_min=0, last_epoch=-1) if apex_support and self.config["fp16_precision"]: model, optimizer = amp.initialize(model, optimizer, opt_level="O2", keep_batchnorm_fp32=True) model_checkpoints_folder = os.path.join(self.writer.log_dir, "checkpoints") # save config file _save_config_file(model_checkpoints_folder) n_iter = 0 valid_n_iter = 0 best_valid_loss = np.inf for epoch_counter in range(self.config["epochs"]): print("====== Epoch {} =======".format(epoch_counter)) for (xis, xjs), _ in train_loader: optimizer.zero_grad() xis = xis.to(self.device) xjs = xjs.to(self.device) loss = self._adv_step(model, augmentor, xis, xjs, n_iter) if n_iter % self.config["log_every_n_steps"] == 0: self.writer.add_scalar("train_loss", loss, global_step=n_iter) if apex_support and self.config["fp16_precision"]: with amp.scale_loss(loss, optimizer) as scaled_loss: scaled_loss.backward() else: loss.backward() # for p in augmentor.parameters(): # # print(p.name) # p.grad *= -1.0 optimizer.step() # Update augmentor augmentor_optimizer.zero_grad() loss = self._adv_step(model, augmentor, xis, xjs, n_iter) if self.augmentor_loss_type == "hinge": loss = torch.clamp(loss, 0.0, 5.4) loss *= -1.0 loss.backward() augmentor_optimizer.step() n_iter += 1 # validate the model if requested if epoch_counter % self.config["eval_every_n_epochs"] == 0: valid_loss = self._validate(model, augmentor, valid_loader) if valid_loss < best_valid_loss: # save the model weights best_valid_loss = valid_loss torch.save( model.state_dict(), os.path.join(model_checkpoints_folder, "model.pth"), ) print("validation loss: ", valid_loss) self.writer.add_scalar("validation_loss", valid_loss, global_step=valid_n_iter) valid_n_iter += 1 # warmup for the first 10 epochs if epoch_counter >= 10: scheduler.step() augmentor_scheduler.step() self.writer.add_scalar("cosine_lr_decay", scheduler.get_lr()[0], global_step=n_iter)
def test(self): train_loader, valid_loader, test_loader = self.dataset.get_data_loaders( ) model = ResNetSimCLR(**self.config["model"]).to(self.device) model = self._load_pre_trained_weights(model) self.eval(test_loader, model)
test_loss = (test_loss * i + criterion(out, labels).item()) / (i + 1) res = torch.argmax(out, 1) correct += (res == labels).sum().item() tot += labels.size(0) return 100 - 100 * correct / tot, test_loss config_path = "./best_run/config.yaml" model_path = "./best_run/model.pth" data_path = "./data" config = yaml.load(open(config_path, "r"), Loader=yaml.FullLoader) device = 'cuda:1' if torch.cuda.is_available() else 'cpu' model = ResNetSimCLR(**config["model"]).to(device) # x = torch.randn(2, 3, 32, 32).to(device) # print(model(x)) # exit() state_dict = torch.load(model_path) model.load_state_dict(state_dict) lr = 0.001 num_epoch = 200 batch_size = 512 num_classes = 10 transform = transforms.Compose( [transforms.ToTensor(),]) # may need to rewrite