# Model and optimizer if args.model == 'AE' or args.model == 'PCA': model = AE(n_feat=n_feat, n_hid=args.hidden, n_lat=args.latent, dropout=args.dropout) elif args.model == 'GCAE': model = GCAE(n_feat=n_feat, n_hid=args.hidden, n_lat=args.latent, dropout=args.dropout) else: raise ValueError("You choose wrong network model") optimizer = optim.Adam(model.parameters(), lr=args.lr) criterion = nn.MSELoss() if args.checkpoint is not None: checkpoint = torch.load(args.checkpoint) model.load_state_dict(checkpoint['model_state_dict']) optimizer.load_state_dict(checkpoint['optimizer_state_dict']) n = len(adj) def train(epoch): t = time.time() running_loss = 0 for i in range(n):
class BiomeAE(): def __init__(self, args): if args.model in ["BiomeAEL0"]: self.mlp_type = "L0" else: self.mlp_type = None self.model_alias = args.model_alias self.model= args.model self.snap_loc = os.path.join(args.vis_dir, "snap.pt") #tl.configure("runs/ds.{}".format(model_alias)) #tl.log_value(model_alias, 0) """ no stat file needed for now stat_alias = 'obj_DataStat+%s_%s' % (args.dataset_name, args.dataset_subset) stat_path = os.path.join( output_dir, '%s.pkl' % (stat_alias) ) with open(stat_path,'rb') as sf: data_stats = pickle.load(sf) """ os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu_num self.predictor = None torch.manual_seed(args.seed) if torch.cuda.is_available(): torch.cuda.manual_seed(args.seed) self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') def get_transformation(self): return None def loss_fn(self, recon_x, x, mean, log_var): if self.model in ["BiomeAE","BiomeAESnip"]: mseloss = torch.nn.MSELoss() return torch.sqrt(mseloss(recon_x, x)) if self.model in ["BiomeAEL0"]: mseloss = torch.nn.MSELoss() return torch.sqrt(mseloss(recon_x, x))+self.predictor.regularization() elif self.model =="BiomeVAE": BCE = torch.nn.functional.binary_cross_entropy( recon_x.view(-1, 28*28), x.view(-1, 28*28), reduction='sum') KLD = -0.5 * torch.sum(1 + log_var - mean.pow(2) - log_var.exp()) return (BCE + KLD) / x.size(0) def param_l0(self): return self.predictor.param_l0() def init_fit(self, X1_train, X2_train, y_train, X1_val, X2_val, y_val, args, ): self.train_loader = get_dataloader (X1_train, X2_train, y_train, args.batch_size) self.test_loader = get_dataloader(X1_val, X2_val, y_val, args.batch_size) self.predictor = AE( encoder_layer_sizes=[X1_train.shape[1]], latent_size=args.latent_size, decoder_layer_sizes=[X2_train.shape[1]], activation=args.activation, batch_norm= args.batch_norm, dropout=args.dropout, mlp_type=self.mlp_type, conditional=args.conditional, num_labels=10 if args.conditional else 0).to(self.device) self.optimizer = torch.optim.Adam(self.predictor.parameters(), lr=args.learning_rate) self.scheduler = torch.optim.lr_scheduler.ExponentialLR(self.optimizer, gamma=0.8) def train(self, args): if args.contr: print("Loading from ", self.snap_loc) loaded_model_para = torch.load(self.snap_loc) self.predictor.load_state_dict(loaded_model_para) t = 0 logs = defaultdict(list) iterations_per_epoch = len(self.train_loader.dataset) / args.batch_size num_iterations = int(iterations_per_epoch * args.epochs) for epoch in range(args.epochs): tracker_epoch = defaultdict(lambda: defaultdict(dict)) for iteration, (x1, x2, y) in enumerate(self.train_loader): t+=1 x1, x2, y = x1.to(self.device), x2.to(self.device), y.to(self.device) if args.conditional: x2_hat, z, mean, log_var = self.predictor(x1, y) else: x2_hat, z, mean, log_var = self.predictor(x1) for i, yi in enumerate(y): id = len(tracker_epoch) tracker_epoch[id]['x'] = z[i, 0].item() tracker_epoch[id]['y'] = z[i, 1].item() tracker_epoch[id]['label'] = yi.item() loss = self.loss_fn(x2_hat, x2, mean, log_var) self.optimizer.zero_grad() loss.backward() if (t + 1) % int(num_iterations / 10) == 0: self.scheduler.step() self.optimizer.step() #enforce non-negative if args.nonneg_weight: for layer in self.predictor.modules(): if isinstance(layer, nn.Linear): layer.weight.data.clamp_(0.0) logs['loss'].append(loss.item()) if iteration % args.print_every == 0 or iteration == len(self.train_loader)-1: print("Epoch {:02d}/{:02d} Batch {:04d}/{:d}, Loss {:9.4f}".format( epoch, args.epochs, iteration, len(self.train_loader)-1, loss.item())) if args.model =="VAE": if args.conditional: c = torch.arange(0, 10).long().unsqueeze(1) x = self.predictor.inference(n=c.size(0), c=c) else: x = self.predictor.inference(n=10) if not args.contr: print("Saving to ", self.snap_loc) torch.save(self.predictor.state_dict(), self.snap_loc) def fit(self,X1_train, X2_train, y_train, X1_val, X2_val, y_val, args,): self.init_fit(X1_train, X2_train, y_train, X1_val, X2_val, y_val, args) self.train(args) def get_graph(self): """ return nodes and weights :return: """ nodes = [] weights = [] for l, layer in enumerate(self.predictor.modules()): if isinstance(layer, nn.Linear): lin_layer =layer nodes.append(["%s"%(x) for x in list(range(lin_layer.in_features))]) weights.append(lin_layer.weight.detach().cpu().numpy().T) nodes.append(["%s"%(x) for x in list(range(lin_layer.out_features))]) #last linear layer return (nodes, weights) def predict(self,X1_val, X2_val, y_val, args): #Batch test x1, x2, y = torch.FloatTensor(X1_val).to(self.device), torch.FloatTensor(X2_val).to(self.device), torch.FloatTensor(y_val).to(self.device) if args.conditional: x2_hat, z, mean, log_var = self.predictor(x1, y) else: x2_hat, z, mean, log_var = self.predictor(x1) val_loss = self.loss_fn( x2_hat, x2, mean, log_var) print("val_loss: {:9.4f}", val_loss.item()) return x2_hat.detach().cpu().numpy() def transform(self,X1_val, X2_val, y_val, args): x1, x2, y = torch.FloatTensor(X1_val).to(self.device), torch.FloatTensor(X2_val).to( self.device), torch.FloatTensor(y_val).to(self.device) if args.conditional: x2_hat, z, mean, log_var = self.predictor(x1, y) else: x2_hat, z, mean, log_var = self.predictor(x1) return z.detach().cpu().numpy() def get_influence_matrix(self): return self.predictor.get_influence_matrix()
up_transform_list = [ utils.to_sparse(up_transform).to(device) for up_transform in tmp['up_transform'] ] model = AE(args.in_channels, args.out_channels, args.latent_channels, edge_index_list, down_transform_list, up_transform_list, K=args.K).to(device) print(model) if args.optimizer == 'Adam': optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) elif args.optimizer == 'SGD': optimizer = torch.optim.SGD(model.parameters(), lr=args.lr, weight_decay=args.weight_decay, momentum=0.9) else: raise RuntimeError('Use optimizers of SGD or Adam') scheduler = torch.optim.lr_scheduler.StepLR(optimizer, args.decay_step, gamma=args.lr_decay) ## -------------- test code start here -------------------- ## Load data from ply file and make it ready
def data_aug(data, lr=0.001, epoch=800, batch_size=128): folder = 'data_aug' save_path = f'{folder}/data_augment.csv' clean_file(save_path) store_e = [700, 750, 800] if not os.path.exists(folder): os.makedirs(folder) else: for i in store_e: result = test(data, folder, i) return result train_loss_curve = [] valid_loss_curve = [] # load model device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model = AE() model = model.to(device) model.train() dataset = AEDataset(data) train_size = int(0.85 * len(dataset)) valid_size = len(dataset) - train_size train_dataset, valid_dataset = random_split(dataset, [train_size, valid_size]) train_dataloader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True) valid_dataloader = DataLoader(dataset=valid_dataset, batch_size=batch_size, shuffle=True) # loss function and optimizer # can change loss function and optimizer you want criterion = nn.MSELoss() optimizer = torch.optim.Adam(model.parameters(), lr=lr) best = 100 # start training for e in range(epoch): train_loss = 0.0 valid_loss = 0.0 print(f'\nEpoch: {e+1}/{epoch}') print('-' * len(f'Epoch: {e+1}/{epoch}')) # tqdm to disply progress bar for inputs in tqdm(train_dataloader): # data from data_loader inputs = inputs.float().to(device) outputs = model(inputs, device) loss = criterion(outputs, inputs) optimizer.zero_grad() loss.backward() optimizer.step() train_loss += loss.item() for inputs in tqdm(valid_dataloader): # data from data_loader inputs = inputs.float().to(device) outputs = model(inputs, device) # MSE loss loss = criterion(outputs, inputs) # loss calculate valid_loss += loss.item() # save the best model weights as .pth file loss_epoch = train_loss / len(train_dataset) valid_loss_epoch = valid_loss / len(valid_dataset) # if valid_loss_epoch < best : # best = valid_loss_epoch # torch.save(model.state_dict(), 'data_aug.pth') if e in store_e: torch.save(model.state_dict(), f'{folder}/ep{e}data_aug.pth') print(f"training in epoch {e},start augment data!!") result = test(data, folder, e) print(f'Training loss: {loss_epoch:.4f}') print(f'Valid loss: {valid_loss_epoch:.4f}') # save loss every epoch train_loss_curve.append(loss_epoch) valid_loss_curve.append(valid_loss_epoch) # generate training curve # visualize(train_loss_curve,valid_loss_curve, 'Data Augmentation') return result
transform=transforms.Compose([ transforms.RandomHorizontalFlip(), transforms.ToTensor() ])) data_loader = torch.utils.data.DataLoader(dataset=data, batch_size=BATCH_SIZE, shuffle=True) ### INIT MODEL device = torch.device("cpu") #change if on GPU, also need to use .cuda() model = AE().to(device) ### MSE LOSS AND ADAM OPTIMIZER criterion = nn.MSELoss(size_average=True, reduce=True) optimizer = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE) ### TRAIN for epoch in range(EPOCHS): for idx, (sample, _) in enumerate(data_loader): if sample.shape[0] < BATCH_SIZE: break reconstruction, z = model(sample) loss = criterion(reconstruction, sample) optimizer.zero_grad() loss.sum().backward() optimizer.step() if idx % LOG_INTERVAL == 0: print( f"Train Epoch: {epoch} [{idx * len(sample)} / {len(data_loader.dataset)}] - Loss: {loss.sum()}" )
# Model and optimizer if args.model == 'AE' or args.model == "PCA": model = AE(n_feat=n_feat, n_hid=args.hidden, n_lat=args.latent, dropout=args.dropout) elif args.model == 'GCAE': model = GCAE(n_feat=n_feat, n_hid=args.hidden, n_lat=args.latent, dropout=args.dropout) else: raise ValueError("You choose wrong network model") optimizer = optim.Adam(model.parameters()) criterion = nn.MSELoss() if args.checkpoint is not None: checkpoint = torch.load(args.checkpoint) model.load_state_dict(checkpoint['model_state_dict']) optimizer.load_state_dict(checkpoint['optimizer_state_dict']) n = len(adj) def test(): if args.model == 'PCA': pca = PCA(n_components=10) x, y = prepare_svm_data(features, labels) x_pca = pca.fit_transform(x)
#Be sure to use abspath linux does not give the path if one uses __file__ _BASEDIR = os.path.dirname(os.path.abspath(__file__)) src_dir = os.path.join(_BASEDIR, 'src') sys.path.append(src_dir) from flow import ROIFlowBatch from models import AE mask_file = '/Volumes/behavgenom_archive$/Avelino/screening/CeNDR/MaskedVideos/CeNDR_Set1_160517/BRC20067_worms10_food1-10_Set10_Pos5_Ch6_16052017_165021.hdf5' feat_file = '/Volumes/behavgenom_archive$/Avelino/screening/CeNDR/Results/CeNDR_Set1_160517/BRC20067_worms10_food1-10_Set10_Pos5_Ch6_16052017_165021_featuresN.hdf5' if __name__ == '__main__': ae = AE(64) criterion = nn.MSELoss() optimizer = torch.optim.Adam(ae.parameters(), lr=1e-2) #%% #optimizer = torch.optim.SGD(vae.parameters(), lr=1e-1, momentum=0.9) n_epochs = 100 #%% gen = ROIFlowBatch(mask_file, feat_file, batch_size=32, roi_size=128) ae.train() for epoch in range(n_epochs): pbar = tqdm.tqdm(gen) for X in pbar: decoded = ae.forward(X) loss = criterion(decoded, X) # mean square error optimizer.zero_grad() # clear gradients for this training step loss.backward() # backpropagation, compute gradients
def main(args): model_alias = 'DeepBiome_%s+%s_%s+fea1_%s+fea2_%s+bs_%s+%s' % ( args.model, args.dataset_name, args.data_type, args.fea1,args.fea2, args.batch_size, args.extra) tl.configure("runs/ds.{}".format(model_alias)) tl.log_value(model_alias, 0) """ no stat file needed for now stat_alias = 'obj_DataStat+%s_%s' % (args.dataset_name, args.dataset_subset) stat_path = os.path.join( output_dir, '%s.pkl' % (stat_alias) ) with open(stat_path,'rb') as sf: data_stats = pickle.load(sf) """ os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu_num data_path = os.path.join(DATA_ROOT, "ibd_{}.pkl".format(args.data_type)) logger.info("Initializing train dataset") # load data print('==> loading data'); print() (X1_train, X2_train, y_train), (X1_val, X2_val, y_val) = load_data(data_path) train_loader = get_dataloader (X1_train, X2_train, y_train, args.batch_size) test_loader = get_dataloader(X1_val, X2_val, y_val, args.batch_size) torch.manual_seed(args.seed) if torch.cuda.is_available(): torch.cuda.manual_seed(args.seed) device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') ts = time.time() def loss_fn(model, recon_x, x, mean, log_var): if model =="AE": mseloss = torch.nn.MSELoss() return torch.sqrt(mseloss(recon_x, x)) elif model =="VAE": BCE = torch.nn.functional.binary_cross_entropy( recon_x.view(-1, 28*28), x.view(-1, 28*28), reduction='sum') KLD = -0.5 * torch.sum(1 + log_var - mean.pow(2) - log_var.exp()) return (BCE + KLD) / x.size(0) if args.model == "AE": predictor = AE( encoder_layer_sizes=[X1_train.shape[1]], latent_size=args.latent_size, decoder_layer_sizes=[X2_train.shape[1]], activation=args.activation, batch_norm= args.batch_norm, dropout=args.dropout, conditional=args.conditional, num_labels=10 if args.conditional else 0).to(device) else: predictor = VAE( encoder_layer_sizes=args.encoder_layer_sizes, latent_size=args.latent_size, decoder_layer_sizes=args.decoder_layer_sizes, activation=args.activation, batch_norm=args.batch_norm, dropout=args.dropout, conditional=args.conditional, num_labels=10 if args.conditional else 0).to(device) optimizer = torch.optim.Adam(predictor.parameters(), lr=args.learning_rate) logs = defaultdict(list) for epoch in range(args.epochs): tracker_epoch = defaultdict(lambda: defaultdict(dict)) for iteration, (x1, x2, y) in enumerate(train_loader): x1, x2, y = x1.to(device), x2.to(device), y.to(device) if args.conditional: x2_hat, z, mean, log_var = predictor(x1, y) else: x2_hat, z, mean, log_var = predictor(x1) for i, yi in enumerate(y): id = len(tracker_epoch) tracker_epoch[id]['x'] = z[i, 0].item() tracker_epoch[id]['y'] = z[i, 1].item() tracker_epoch[id]['label'] = yi.item() loss = loss_fn(args.model, x2_hat, x2, mean, log_var) optimizer.zero_grad() loss.backward() optimizer.step() logs['loss'].append(loss.item()) if iteration % args.print_every == 0 or iteration == len(train_loader)-1: print("Epoch {:02d}/{:02d} Batch {:04d}/{:d}, Loss {:9.4f}".format( epoch, args.epochs, iteration, len(train_loader)-1, loss.item())) if args.model =="VAE": if args.conditional: c = torch.arange(0, 10).long().unsqueeze(1) x = predictor.inference(n=c.size(0), c=c) else: x = predictor.inference(n=10) """ plt.figure() plt.figure(figsize=(5, 10)) for p in range(10): plt.subplot(5, 2, p+1) if args.conditional: plt.text( 0, 0, "c={:d}".format(c[p].item()), color='black', backgroundcolor='white', fontsize=8) plt.imshow(x[p].view(28, 28).cpu().data.numpy()) plt.axis('off') if not os.path.exists(os.path.join(args.fig_root, str(ts))): if not(os.path.exists(os.path.join(args.fig_root))): os.mkdir(os.path.join(args.fig_root)) os.mkdir(os.path.join(args.fig_root, str(ts))) plt.savefig( os.path.join(args.fig_root, str(ts), "E{:d}I{:d}.png".format(epoch, iteration)), dpi=300) plt.clf() plt.close('all') """ #Batch test x1, x2, y = torch.FloatTensor(X1_val).to(device), torch.FloatTensor(X2_val).to(device), torch.FloatTensor(y_val).to(device) if args.conditional: x2_hat, z, mean, log_var = predictor(x1, y) else: x2_hat, z, mean, log_var = predictor(x1) val_loss = loss_fn(args.model, x2_hat, x2, mean, log_var) print("val_loss: {:9.4f}", val_loss.item()) """