def train(): graphs, labels = load_data("datasets/train") train_inputs, train_targets, val_inputs, val_targets = utils.split_train_val( graphs, labels, val_rate=0.3) model = GNNModel(8) loss_func = BinaryCrossEntropy() optimizer = Adam() batch_generator = utils.BatchGenerator(batch_size=32) min_loss = 100000 for epoch in range(50): print(f"Epoch{epoch + 1}") train_losses = [] for inputs, targets in batch_generator.generator( train_inputs, train_targets): train_loss, loss_grad = loss_func(model, inputs, targets, is_grad=True) optimizer.update(model, loss_grad) train_losses.append(train_loss) train_mean_loss = np.mean(train_losses) pred = np.array([model.predict(input_) for input_ in train_inputs]).squeeze() train_accuracy = accuracy(pred, train_targets) val_losses = [] for inputs, targets in batch_generator.generator( val_inputs, val_targets): val_loss, _ = loss_func(model, inputs, targets, is_grad=False) val_losses.append(val_loss) val_mean_loss = np.mean(val_losses) pred = np.array([model.predict(input_) for input_ in val_inputs]).squeeze() val_accuracy = accuracy(pred, val_targets) if min(min_loss, val_mean_loss) < min_loss: min_loss = val_mean_loss print( f"Train loss: {train_mean_loss}\tTrain accuracy: {train_accuracy}" ) print( f"Validation loss: {val_mean_loss}\tValidation accuracy: {val_accuracy}" ) print("")
def init(config): ''' Loads the GloVe embeddings for the words which occur in the IMDB train set vocab and uses that vocab to create train, validation and test sets for the IMDB dataset. Extracts the pad_id token. ''' import os if not os.path.isdir('.data'): os.mkdir('.data') # Extract the initial vocab from the IMDB dataset vocab = IMDB(data_select='train')[0].get_vocab() # Create GloVe embeddings based on original vocab # word freqs glove_vocab = torchtext.vocab.Vocab( counter=vocab.freqs, max_size=MAX_VOCAB_SIZE, min_freq=MIN_FREQ, vectors=torchtext.vocab.GloVe(name='6B')) # Acquire 'Spacy' tokenizer for the vocab words tokenizer = get_tokenizer('spacy', 'en_core_web_sm') # Acquire train and test IMDB sets with previously created # GloVe vocab and 'Spacy' tokenizer train_set, test_set = IMDB(tokenizer=tokenizer, vocab=glove_vocab) # Extract the vocab of the acquired train set vocab = train_set.get_vocab() # Extract the token used for padding pad_id = vocab['<pad>'] # Split the train set into train and validation sets train_set, valid_set = split_train_val(train_set) config['train'] = train_set config['val'] = valid_set config['test'] = test_set config['vocab'] = vocab config['pad_id'] = pad_id
if __name__ == '__main__': # 命令行参数 cfg_filename = sys.argv[1] # 配置文件名 mode = sys.argv[2] # 运行模式,包括train,test,online_test # 读取配置文件 cfg = Config.fromfile(cfg_filename) print("config filename: " + str(cfg_filename)) set_seed(cfg.random_seed) # 划分数据集 if mode == 'split_data': split_train_val() # 训练模型 if mode == 'train': train_main(cfg=cfg) # 测试模型 if mode == 'test': test_main(cfg=cfg) # 复赛线上测试 if mode == 'test_online': test_online_main(cfg=cfg) print('end')
# plt.plot(hist_coal, 'r-') # plt.plot(hist_gangue, 'g-') # plt.xlim([0, 256]) # plt.show() from sklearn.metrics import precision_recall_curve, roc_curve, classification_report from utils import get_imgs_and_masks, get_ids, split_train_val from unet import UNet import torch ori_w, ori_h = 852, 480 dir_img = '/home/zhuzhu/Desktop/mid project/raw_data' dir_mask = '/home/zhuzhu/Desktop/mid project/groundtruth' ids = get_ids(dir_img) iddataset = split_train_val(ids, 0.05) net = UNet(1, 2) net.eval() net.load_state_dict( torch.load( '/media/zhuzhu/0C5809B80C5809B8/draft/unet/checkpoint/unet_0.854608765.pth', map_location='cpu')) val = get_imgs_and_masks(iddataset['val'], dir_img, dir_mask) c = 0 for i, b in enumerate(val): img = np.array(b[0]).astype(np.float32) mask = np.array(b[1]).astype(np.float32) with torch.no_grad():
def train_net(net, epochs=5, batch_size=1, lr=0.1, val_percent=0.05, save_cp=True, gpu=False, img_scale=0.5, dir_img=None, dir_mask=None, dir_checkpoint=None, channels=1, classes=1): ids = os.listdir(dir_img) if not os.path.exists(dir_checkpoint): os.makedirs(dir_checkpoint, mode=0o755) iddataset = split_train_val(ids, val_percent) print('Starting training:') print('Epochs: ' + str(epochs)) print('Batch size: ' + str(batch_size)) print('Learning rate: ' + str(lr)) print('Training size: ' + str(len(iddataset['train']))) print('Validation size: ' + str(len(iddataset['val']))) print('Checkpoints: ' + str(save_cp)) N_train = len(iddataset['train']) optimizer = optim.RMSprop(net.parameters(), lr=lr) criterion = nn.BCELoss() for epoch in range(epochs): print('Starting epoch {}/{}.'.format(epoch + 1, epochs)) net.train() # reset the generators train = get_imgs_and_masks(iddataset['train'], dir_img, dir_mask, img_scale) val = get_imgs_and_masks(iddataset['val'], dir_img, dir_mask, img_scale) epoch_loss = 0 # Run Batch for i, b in enumerate(batch(train, batch_size)): # Grab data try: imgs = np.array([i[0] for i in b]).astype(np.float32) true_masks = np.array([i[1] for i in b]) except: print( 'prob have dimension issues, wrong orientations or half reconned images' ) # Deal with dimension issues if channels == 1: imgs = np.expand_dims(imgs, 1) if classes > 1: true_masks = to_categorical(true_masks, num_classes=classes) # Play in torch's sandbox imgs = torch.from_numpy(imgs) true_masks = torch.from_numpy(true_masks) # Send to GPU if gpu: imgs = imgs.cuda() true_masks = true_masks.cuda() # Predicted segmentations masks_pred = net(imgs) # Flatten masks_probs_flat = masks_pred.view(-1) true_masks_flat = true_masks.view(-1) # Calculate losses btwn true/predicted loss = criterion(masks_probs_flat, true_masks_flat) epoch_loss += loss.item() # Batch Loss print('{0:.4f} --- loss: {1:.6f}'.format(i * batch_size / N_train, loss.item())) # Backprop optimizer.zero_grad() loss.backward() optimizer.step() # Epoch Loss print('Epoch finished ! Loss: {}'.format(epoch_loss / i)) if 1: val_dice = eval_net(net, val, epoch, dir_checkpoint, gpu) print('Validation Dice Coeff: {}'.format(val_dice)) if save_cp: torch.save( net.state_dict(), os.path.join(dir_checkpoint, 'CP{}.pth'.format(epoch + 1))) print('Checkpoint {} saved !'.format(epoch + 1))
def train_net(net, epochs=5, batch_size=1, lr=1e-3, val_percent=0.05, save_cp=True, gpu=False, img_scale=0.5): dir_img = '/home/xyj/data/spacenet/vegas/images_rgb_1300/' dir_mask = '/home/xyj/test/Pytorch-UNet/data/train_mask_point/' dir_checkpoint = 'checkpoints_point/' if not os.path.exists(dir_checkpoint): os.mkdir(dir_checkpoint) # ids = get_ids(dir_img) # 返回train文件夹下文件的名字列表,生成器(except last 4 character,.jpg这样的) with open('train_list.txt', 'r') as f: lines = f.readlines() ids = (i.strip('\n')[:-4] for i in lines) ids = split_ids( ids) # 返回(id, i), id属于ids,i属于range(n),相当于把train的图✖️了n倍多张,是tuple的生成器 iddataset = split_train_val( ids, val_percent ) # validation percentage,是dict = {"train": ___(一个list), "val":___(一个list)} print(''' Starting training: Epochs: {} Batch size: {} Learning rate: {} Training size: {} Validation size: {} Checkpoints: {} CUDA: {} '''.format(epochs, batch_size, lr, len(iddataset['train']), len(iddataset['val']), str(save_cp), str(gpu))) N_train = len(iddataset['train']) # optimizer = optim.SGD(net.parameters(), # lr=lr, # momentum=0.9, # weight_decay=0.0005) optimizer = optim.Adam(net.parameters(), lr=lr, betas=(0.9, 0.999), eps=1e-3) # scheduler = optim.lr_scheduler.StepLR(optimizer,step_size=40,gamma = 0.3) criterion = nn.BCELoss() for epoch in range(epochs): print('Starting epoch {}/{}.'.format(epoch + 1, epochs)) net.train() # reset the generators train = get_imgs_and_masks(iddataset['train'], dir_img, dir_mask, img_scale) val = get_imgs_and_masks(iddataset['val'], dir_img, dir_mask, img_scale) epoch_loss = 0 for i, b in enumerate(batch(train, batch_size)): imgs = np.array([i[0] for i in b]).astype(np.float32) true_masks = np.array([i[1] // 200 for i in b]) imgs = torch.from_numpy(imgs) true_masks = torch.from_numpy(true_masks) if gpu: imgs = imgs.cuda() true_masks = true_masks.cuda() masks_pred = net(imgs) masks_probs_flat = masks_pred.view(-1) true_masks_flat = true_masks.view(-1) loss = criterion(masks_probs_flat, true_masks_flat) epoch_loss += loss.item() print('{0:.4f} --- loss: {1:.6f}'.format(i * batch_size / N_train, loss.item())) optimizer.zero_grad() loss.backward() optimizer.step() # scheduler.step() print('Epoch finished ! Loss: {}'.format(epoch_loss / i)) if 1: val_dice = eval_net(net, val, gpu) print('Validation Dice Coeff: {}'.format(val_dice)) if save_cp: torch.save(net.state_dict(), dir_checkpoint + 'CP{}.pth'.format(epoch + 1)) print('Checkpoint {} saved !'.format(epoch + 1))
for img in list(tif.iter_images()): tiff = img cv2.imwrite(os.path.join(dir_path, name[0] + '.jpg')) img_path = os.path.join(dir_path, name[0] + '.jpg') xml_path = os.path.join(dir_path, name[0] + '.xml') img_table, fm_table = preprocess(img_path, xml_path) assert len(img_table) == len( fm_table), "image and feature map not match" for idx in range(len(img_table)): image_name.append(name[0] + '_' + str(idx)) cv2.imwrite( os.path.join(save_path, name[0] + '_' + str(idx) + '.jpg'), img_table[idx]) cv2.imwrite( os.path.join(save_path, name[0] + '_' + str(idx) + '_maskh.png'), fm_table[idx][0]) cv2.imwrite( os.path.join(save_path, name[0] + '_' + str(idx) + '_maskv.png'), fm_table[idx][1]) dir_path = r'D:\a上财学习材料\20毕业论文\实验\data' filename_path = 'cTDaR19_imagename.txt' write_filename(os.path.join(dir_path, filename_path), image_name) # split train and test split_train_val(dir_path, filename_path)
def train_net(net, epochs=5, batch_size=1, lr=0.003, val_percent=0.20, loss_lambda=5, save_cp=True, gpu=False, img_scale=0.5, expositions_num=15, logg_freq=15, tb=False, w_decay=0.0005, use_notifications=False, polyaxon=False, outputs_path='checkpoints'): # === Localize training data =================================================== if polyaxon: data_paths = get_data_paths() dir_checkpoints = get_outputs_path() dataSets_dir = os.path.join(data_paths['data1'], 'eprado', 'USLDR-DataSet') #dataSets_dir = os.path.join(data_paths['data1'] , 'eprado', 'LDR_DataSet') else: dataSets_dir = os.path.join(wk_dir, "LDR_DataSet") dir_checkpoints = os.path.join(wk_dir, outputs_path) print('Dataset_dir', dataSets_dir) print('Outputs_path', dir_checkpoints) experiment_id = datetime.datetime.now().strftime('%d%m_%H%M_') experiment_name = 'ExpandnetL_psn_{}_bs{}_lr{}_exps{}'.format( experiment_id, batch_size, lr, expositions_num) dir_img = os.path.join(dataSets_dir, 'Org_images/') dir_compressions = os.path.join(dataSets_dir, 'c_images/') dir_mask = os.path.join(dataSets_dir, 'c_images/') #if tb: #dummy_input = torch.rand(1, 3, 128, 128) #writer.add_graph(net, (dummy_input,)) #writer.close() # === Load Training/Validation data ===================================================== ids = get_ids(dir_compressions) # Split into train test idsset = list(ids) kf = KFold(n_splits=5, shuffle=False) #print('Train splits: ',kf.get_n_splits(dataset)) best_psnr_m = 0 best_psnr_hvs = 0 #for train_index, test_index in kf.split(idsset): iddataset = split_train_val(idsset, expositions_num, val_percent) #test_set = [] #for im_id in test_index: # for e in range(expositions_num): # test_set.append(idsset[im_id]) N_train = len(iddataset['train']) N_val = len(iddataset['val']) N_test = 0 #len(test_set) #=====CHOOSE Loss Criterion============================================================= #criterion = nn.MSELoss(reduction='mean') criterion = ExpandNetLoss(loss_lambda=loss_lambda) optimizer = optim.Adagrad(net.parameters(), lr=lr, lr_decay=0.000001, weight_decay=w_decay) #optimizer = optim.SGD(net.parameters(), # lr=lr, # momentum=0.9, # weight_decay=0.0005) since = time.time() print(''' Training SETUP: Epochs: {0:} Batch size: {1:} Optimizer: Adagrad Learning rate: {2:} Weight decay: {3:} Training size: {4:} Validation size: {5:} Test size: {6:} Checkpoints: {7:} CUDA: {8:} '''.format(epochs, batch_size, lr, w_decay, N_train, N_val, N_test, str(save_cp), str(gpu))) train_dataset = HdrDataset(iddataset['train'], dir_compressions, dir_mask, expositions_num) val_dataset = HdrDataset(iddataset['val'], dir_compressions, dir_mask, expositions_num) #test_dataset = HdrDataset(test_set, dir_compressions, dir_mask,expositions_num) train_data_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, drop_last=False) val_data_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, drop_last=False) #test_data_loader = DataLoader(test_dataset,batch_size=batch_size,shuffle=True) best_hvsm = 0.0 global_psnr_m = [] global_psnr_hvs = [] for epoch in range(epochs): print('\n') print('{}{}{}'.format('+', '=' * 78, '+')) print('| Starting epoch {}/{}. {}'.format(epoch + 1, epochs, (' ' * 57) + '|')) print('{}{}{}'.format('|', '-' * 78, '|')) begin_of_epoch = time.time() tot_steps = math.trunc(N_train / batch_size) net.train() train_loss = 0 losses = [] val_loss = 0 step = 0 train_sample = [] train_acc = 0 val_hvsm = 0 val_hvs = 0 model_pnsr_m = 0 for i, b in enumerate(train_data_loader): step += 1 imgs, true_masks, imgs_ids = b['input'], b['target'], b['id'] #print(i, b['input'].size(), b['target'].size()) #input: [15, 3, 224, 224]), target: [15, 3, 224, 224] #print('>>>>>>> Input max: ' , torch.max(imgs[0])) #print('>>>>>>> mask max : ', torch.max(true_masks[0])) if gpu: imgs = imgs.cuda() true_masks = true_masks.cuda() else: print(' GPU not available') # Predicted mask images optimizer.zero_grad() prediction = net(imgs) #prediction shape: [B, 3, 224, 224] #cost, cost_input_output = Hdr_loss(imgs, true_masks, prediction, sep_loss=False, gpu=gpu, tb=tb) cost = criterion(prediction, true_masks) #loss is torch tensor losses.append(cost.item()) train_loss = np.mean(losses) cost.backward() optimizer.step() if step == 1 or step % logg_freq == 0: #print('| Step: {0:}, cost:{1:}, Train Loss:{2:.9f}, Train Acc:{3:.9f}'.format(step,cost, train_loss,train_acc/step)) print('| Step: {0:}, cost:{1:}, Train Loss:{2:.9f}'.format( step, cost, train_loss)) #Last Step of this Epoch if step == math.trunc(tot_steps): num_in_batch = random.randrange(imgs.size(0)) train_sample_name = imgs_ids[num_in_batch] train_sample = [ imgs[num_in_batch], true_masks[num_in_batch], prediction[num_in_batch] ] t_exp_name = 'Train_' + experiment_name saveTocheckpoint(dir_checkpoints, t_exp_name, train_sample_name, epoch, train_sample[0], train_sample[1], train_sample[2]) if tb: print( '| saving train step {0:} sample : input,target & pred' .format(step)) grid = torchvision.utils.make_grid(train_sample, nrow=3) writer.add_image('train_sample', grid, 0) #if epoch == 1 or epoch % 15 == 0 or epoch == epochs: val_loss, val_hvsm, val_hvs = eval_hdr_net(net, dir_checkpoints, experiment_name, val_data_loader, criterion, epoch, gpu, batch_size, expositions_num=15, tb=tb) if tb: writer.add_scalar('training_loss: ', train_loss, epoch) writer.add_scalar('validation_loss', val_loss, epoch) writer.add_scalar('val_hvsm', val_hvsm, epoch) writer.add_scalar('val_hvs', val_hvs, epoch) writer.add_scalars('losses', { 'training_loss': train_loss, 'val_loss': val_loss }, epoch) if polyaxon: experiment.log_metrics(step=epoch, training_loss=train_loss, validation_loss=val_loss, val_hvsm=val_hvsm, val_hvs=val_hvs) print('{}{}{}'.format('+', '=' * 78, '+')) print('| {0:} Epoch {1:} finished ! {2:}|'.format( ' ' * 28, (epoch + 1), ' ' * 29)) print('{}{}{}'.format('+', '-' * 78, '+')) print('| Summary: Train Loss: {0:0.07}, Val Loss:{1:}'.format( train_loss, val_loss)) print('| Avrg psnr-hvs_m :{0:0.04},Avrg psnr-hvs :{1:0.04}'. format(val_hvsm, val_hvs)) time_epoch = time.time() - begin_of_epoch print('| Epoch ETC: {:.0f}m {:.0f}s'.format(time_epoch // 60, time_epoch % 60)) print('{}{}{}'.format('+', '=' * 78, '+')) if save_cp and (val_hvsm > best_hvsm): best_hvsm = val_hvsm model_path = os.path.join(dir_checkpoints, 'BestCP.pth') torch.save(net.state_dict(), model_path) print('Checkpoint saved !') global_psnr_hvs.append(val_hvs) global_psnr_m.append(val_hvsm) ''' test_psnr_m, test_psnr_hvs = test_hdr_net(model_path,dir_checkpoints, experiment_name, test_data_loader, criterion,gpu,tb) if save_cp and (test_psnr_m > best_psnr_m): best_psnr_m = test_psnr_m best_model_path = os.path.join(dir_checkpoints, 'Best_CP.pth') torch.save(net.state_dict(),best_model_path) print('Best model saved !') ''' print('>' * 80) time_elapsed = time.time() - since print('Training complete in {:.0f}m {:.0f}s'.format( time_elapsed // 60, time_elapsed % 60)) print('Final Average psnr-hvs_m: {:.0f}, psnr-hvs: {:.0f}'.format( np.mean(global_psnr_m), np.mean(global_psnr_hvs))) if tb: writer.close() if use_notifications: end_msg = "train.py finished at: {}(".format( str(datetime.datetime.now())) push = pb.push_note("usHDR: Finish", end_msg)
def train_net(net, device, epochs=5, batch_size=1, lr=0.1, val_percent=0.15, save_cp=True, img_scale=0.5): ids = get_ids(dir_img) iddataset = split_train_val(ids, val_percent) logging.info('''Starting training: Epochs: {epochs} Batch size: {batch_size} Learning rate: {lr} Training size: {len(iddataset["train"])} Validation size: {len(iddataset["val"])} Checkpoints: {save_cp} Device: {device.type} Images scaling: {img_scale} ''') n_train = len(iddataset['train']) n_val = len(iddataset['val']) optimizer = optim.Adam(net.parameters(), lr=lr) #optimizer = optim.SGD(net.parameters(), lr=lr, momentum=0.75) criterion = nn.BCELoss() for epoch in range(epochs): net.train() # reset the generators train = get_imgs_and_masks(iddataset['train'], dir_img, dir_mask, img_scale) val = get_imgs_and_masks(iddataset['val'], dir_img, dir_mask, img_scale) epoch_loss = 0 with tqdm(total=n_train, desc='Epoch {epoch + 1}/{epochs}', unit='img') as pbar: for i, b in enumerate(batch(train, batch_size)): imgs = np.array([i[0] for i in b]).astype(np.float32) true_masks = np.array([i[1] for i in b]) imgs = torch.from_numpy(imgs) true_masks = torch.from_numpy(true_masks) imgs = imgs.to(device=device) true_masks = true_masks.to(device=device) masks_pred = net(imgs) loss = criterion(masks_pred, true_masks) epoch_loss += loss.item() pbar.set_postfix(**{'loss (batch)': loss.item()}) optimizer.zero_grad() loss.backward() optimizer.step() pbar.update(batch_size) if save_cp: try: os.mkdir(dir_checkpoint) logging.info('Created checkpoint directory') except OSError: pass torch.save(net.state_dict(), dir_checkpoint + 'gpu_3.pth') logging.info('gpu_3 saved !') val_dice = eval_net(net, val, device, n_val) logging.info('Validation Dice Coeff: {}'.format(val_dice))
def train_net(net, epochs=5, batch_size=1, lr=0.1, val_percent=0.2, save_cp=True, gpu=False, img_scale=0.5): dir_img = '/media/workspace/DATASET/Eyes/STARE/stare-images' dir_mask = '/media/workspace/DATASET/Eyes/STARE/labels' dir_checkpoint = 'checkpoints/' ids = get_ids(dir_img) ids = split_ids(ids, 4) iddataset = split_train_val(ids, val_percent) print(''' Starting training: Epochs: {} Batch size: {} Learning rate: {} Training size: {} Validation size: {} Checkpoints: {} CUDA: {} '''.format(epochs, batch_size, lr, len(iddataset['train']), len(iddataset['val']), str(save_cp), str(gpu))) N_train = len(iddataset['train']) optimizer = optim.SGD(net.parameters(), lr=lr, momentum=0.9, weight_decay=0.0005) criterion = nn.BCELoss() for epoch in range(epochs): print('Starting epoch {}/{}.'.format(epoch + 1, epochs)) net.train() # reset the generators train = get_imgs_and_masks(iddataset['train'], dir_img, dir_mask, img_scale) val = get_imgs_and_masks(iddataset['val'], dir_img, dir_mask, img_scale) epoch_loss = 0 for i, b in enumerate(batch(train, batch_size)): imgs = np.array([w[0] for w in b]).astype(np.float32) true_masks = np.array([w[1] for w in b]) true_masks[true_masks == 255] = 1 imgs = torch.from_numpy(imgs) true_masks = torch.from_numpy(true_masks) if gpu: imgs = imgs.cuda() true_masks = true_masks.cuda() masks_pred = net(imgs) masks_probs_flat = masks_pred.view(-1) true_masks_flat = true_masks.view(-1) loss = criterion(masks_probs_flat, true_masks_flat) epoch_loss += loss.item() print('{0:.4f} --- loss: {1:.6f}'.format(i * batch_size / N_train, loss.item())) optimizer.zero_grad() loss.backward() optimizer.step() print('Epoch finished ! Loss: {}'.format(epoch_loss / i)) if 1: val_dice = eval_net(net, val, gpu) print('Validation Dice Coeff: {}'.format(val_dice)) if save_cp: if not os.path.exists(dir_checkpoint): os.mkdir(dir_checkpoint) torch.save(net.state_dict(), dir_checkpoint + 'CP{}.pth'.format(epoch + 1)) print('Checkpoint {} saved !'.format(epoch + 1))
def train_net(net, epochs=5, batch_size=1, lr=0.1, val_percent=0.05, save_cp=True, gpu=False, img_height=512, img_scale=0.5): #dir_img = 'carvana-image-masking-challenge/train/' #dir_mask = 'carvana-image-masking-challenge/train_masks/' dir_img = '/root/ctw/train_images_preprocess_other/' dir_mask = '/root/ctw/train_images_mask_preprocess_other/' #dir_img = '/root/ctw/val_images_preprocess_test/' #dir_mask = '/root/ctw/val_images_mask_preprocess_test/' dir_checkpoint = 'checkpoints/' ids = list(get_ids(dir_img)) ids = split_ids(ids) iddataset = split_train_val(ids, val_percent) print(''' Starting training: Epochs: {} Batch size: {} Learning rate: {} Training size: {} Validation size: {} Checkpoints: {} CUDA: {} '''.format(epochs, batch_size, lr, len(iddataset['train']), len(iddataset['val']), str(save_cp), str(gpu))) N_train = len(iddataset['train']) optimizer = optim.Adam(net.parameters(),lr=lr) # optimizer = optim.SGD(net.parameters(), # lr=lr, # momentum=0.92, # weight_decay=0.0005) criterion = nn.BCELoss() #criterion = nn.MSELoss() #import scipy.misc iteration = 0 for epoch in range(epochs): print('Starting epoch {}/{}.'.format(epoch + 1, epochs)) net.train() # reset the generators train = get_imgs_and_masks(iddataset['train'], dir_img, dir_mask, img_height, img_scale) val = get_imgs_and_masks(iddataset['val'], dir_img, dir_mask, img_height, img_scale) epoch_loss = 0 for i, b in enumerate(batch(train, batch_size)): #print(i, len(b)) """ for j in b: #print(j[0].shape, j[1].shape) #print(j[1]) #scipy.misc.toimage(j[0], cmin=0.0, cmax=1.0).save('%s_outfile.jpg'%count) #scipy.misc.toimage(j[1], cmin=0.0, cmax=1.0).save('%s_outmask.jpg'%count) count += 1 """ iteration += 1 try: imgs = np.array([i[0] for i in b]).astype(np.float32) true_masks = np.array([i[1] for i in b]) # print("\nImgs : \n{}".format(np.unique(imgs))) # print("\ntrue mask \n {} ".format(np.unique(true_masks))) #print('%s'%(datetime.datetime.now()), '{0:.4f}'.format(i * batch_size)) except Exception as e: print(Exception) continue imgs = torch.from_numpy(imgs) true_masks = torch.from_numpy(true_masks) if gpu: imgs = imgs.cuda() true_masks = true_masks.cuda() masks_pred = net(imgs) masks_probs_flat = masks_pred.view(-1) true_masks_flat = true_masks.view(-1) loss = criterion(masks_probs_flat, true_masks_flat) epoch_loss += loss.item() if iteration%100==0: print('iter %s'%iteration, '%s'%(datetime.datetime.now()), '{0:.4f} --- loss: {1:.6f}'.format(i * batch_size / N_train, loss.item())) optimizer.zero_grad() loss.backward() optimizer.step() print('Epoch finished ! Loss: {}'.format(epoch_loss / i)) if 1: val_dice = eval_net(net, val, gpu) val_iou=val_dice/(2-val_dice) print('Validation Dice Coeff: {}'.format(val_dice)) print('Validation iouScore : {}'.format(val_iou)) if save_cp: torch.save(net.state_dict(), dir_checkpoint + 'CP{}.pth'.format(epoch + 1)) print('Checkpoint {} saved !'.format(epoch + 1))
def train_primary(model_name, df_name, max_len, train_ratio, text_feature, embeddings_version, embeddings_path, config_primary, models_path, models_df): device = utils.find_device() loss_function, labels_num, dropout, epochs_num, threshold, \ hid_dim_lstm, lin_output_dim, lr, batch_size, momentum = \ utils.read_config_primary(config_primary, "train_primary") train_df, val_df = utils.split_train_val(df_name, train_ratio) train_dataloader, val_dataloader = utils.create_dataloaders_train( train_df, val_df, text_feature, embeddings_version, max_len, batch_size) model = get_model(model_name, embeddings_path, hid_dim_lstm, loss_function, labels_num, dropout, lin_output_dim) model.to(device) optimizer = AdamW(model.parameters(), lr=2e-5, eps=1e-8) train_len = len(train_dataloader) val_len = len(val_dataloader) train_accuracy_list, train_loss_list, accuracy_list_val, loss_list_val = [], [], [], [] best_val_acc = 0.0 total_steps = len(train_dataloader) * epochs_num scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=0, num_training_steps=total_steps) for epoch in range(epochs_num): utils.print_epochs_progress(epoch, epochs_num) start_train = time.clock() acc_num_train = 0.0 loss_scalar = 0.0 model.train() optimizer.zero_grad() t0 = time.time() predictions_dict_train = dict() for batch_idx, batch in enumerate(train_dataloader): if batch_idx % 10 == 0 and not batch_idx == 0: utils.print_batches_progress(t0, batch_idx, train_dataloader) input_ids = batch[0].to(device, dtype=torch.long) masks = batch[1].to(device, dtype=torch.long) labels = batch[2].to(device, dtype=torch.long) hadm_id = batch[3] model.zero_grad() loss, predictions, probabilities = model(input_ids, masks, labels) loss_scalar += loss.item() loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) acc_num_train += utils.add_accuracy(predictions, labels) optimizer.step() scheduler.step() predictions_dict_train = utils.update_predictions_dict( predictions_dict_train, hadm_id, predictions, probabilities, labels) loss_scalar /= train_len end_train = time.clock() total_time = end_train - start_train utils.print_train_epoch(epoch, acc_num_train, train_len, loss_scalar, total_time) train_loss_list.append(round(float(loss_scalar), 3)) train_accuracy_list.append(round(float(acc_num_train / train_len), 3)) utils.print_train_epoch_end(t0) val_acc, val_loss, predictions_dict_val = evaluate_val( model, val_dataloader, device, val_len) accuracy_list_val.append(val_acc) loss_list_val.append(val_loss) if val_acc > best_val_acc: torch.save(model.state_dict(), models_path + 'primary.pkl') best_val_acc = val_acc utils.save_predictions_to_df(predictions_dict_train, models_path, 'train', 'primary') utils.save_predictions_to_df(predictions_dict_val, models_path, 'validation', 'primary') utils.save_model(models_df, 'primary', hid_dim_lstm, labels_num, loss_function, dropout, lin_output_dim, lr, epochs_num, batch_size, momentum, best_val_acc) utils.print_summary(models_path, 'primary', accuracy_list_val) test_acc, predictions_dict_test = evaluate_test(model, device, embeddings_version, max_len, text_feature, batch_size) utils.save_predictions_to_df(predictions_dict_test, models_path, 'test', 'primary') utils.print_test_results(models_path, 'primary', test_acc) return models_df
if __name__ == "__main__": # Data Loader train_datagen = ImageDataGenerator(rotation_range=45, width_shift_range=0.2, height_shift_range=0.2, shear_range=0.2, zoom_range=0.25, horizontal_flip=True, fill_mode='nearest') test_datagen = ImageDataGenerator() # 切分数据 (train_idx, valid_idx, ytr, yv, labels, selected_breed_list) = split_train_val(labels, NUM_CLASSES, seed=SEED) INPUT_SIZE = 299 # width/height of image in pixels (as expected by Xception model) x_train = np.zeros((len(labels), INPUT_SIZE, INPUT_SIZE, 3), dtype='float32') for i, img_id in tqdm(enumerate(labels['id'])): img = read_img(img_id, data_dir, 'train', (INPUT_SIZE, INPUT_SIZE)) x = xception.preprocess_input(np.expand_dims(img.copy(), axis=0)) x_train[i] = x print('\nTotal Images shape: {}'.format(x_train.shape)) Xtr = x_train[train_idx] Xv = x_train[valid_idx] print('Train (images, H, W, C):', Xtr.shape, '\nVal (images, H, W, C):',
def train_net(net, epochs=5, batch_size=1, lr=0.1, val_percent=0.05, save_cp=True, gpu=False, img_scale=1): dir_img = 'data/train/' dir_mask = 'data/train_masks/' dir_checkpoint = 'checkpoints/' ids = get_ids(dir_img) iddataset = split_train_val(ids, val_percent) traindata = BrainDataset(dir_img, dir_mask, iddataset['train']) valdata = BrainDataset(dir_img, dir_mask, iddataset['val']) trainloader = DataLoader(traindata, batch_size=args.batchsize, shuffle=True, num_workers=4) valloader = DataLoader(valdata, batch_size=1, shuffle=False, num_workers=2) print(''' Starting training: Epochs: {} Batch size: {} Learning rate: {} Training size: {} Validation size: {} Checkpoints: {} CUDA: {} '''.format(epochs, batch_size, lr, len(iddataset['train']), len(iddataset['val']), str(save_cp), str(gpu))) N_train = len(iddataset['train']) optimizer = optim.SGD(net.parameters(), lr=lr, momentum=0.9, weight_decay=0.0005) criterion = nn.BCELoss() epoch_tb = 0 # for tensorboard plotting for epoch in range(epochs): print('Starting epoch {}/{}.'.format(epoch + 1, epochs)) if epoch % 4 == 0: for param_group in optimizer.param_groups: param_group['lr'] /= 2 print('learning rate == ' + str(param_group['lr'])) epoch_loss = 0 for i, data in enumerate(trainloader): imgs = data[0] true_masks = data[1] if gpu: imgs = imgs.cuda() true_masks = true_masks.cuda() masks_pred = net(imgs) masks_probs = F.sigmoid(masks_pred) masks_probs_flat = masks_probs.view(-1) true_masks_flat = true_masks.view(-1) true_masks_flat = true_masks_flat.float() loss = criterion(masks_probs_flat, true_masks_flat) epoch_loss += loss.item() print('{0:.4f} --- loss: {1:.6f}'.format(i * batch_size / N_train, loss.item())) writer.add_scalar('loss - iter', loss.item(), epoch_tb) epoch_tb += 1 optimizer.zero_grad() loss.backward() optimizer.step() print('Epoch finished ! Loss: {}'.format(epoch_loss / (i + 1))) writer.add_scalar('epoch_loss - epoch', epoch_loss / (i + 1), epoch) if epoch % 7 == 0: val_dice = eval_net(net, valloader, gpu, threshold=0.3) print('Validation Dice Coeff 0.3: {}'.format(val_dice)) val_dice = eval_net(net, valloader, gpu, threshold=0.4) print('Validation Dice Coeff 0.4: {}'.format(val_dice)) val_dice = eval_net(net, valloader, gpu, threshold=0.5) print('Validation Dice Coeff 0.5: {}'.format(val_dice)) writer.add_scalar('val_dice_coeff - epoch', val_dice, epoch) if save_cp: torch.save(net.state_dict(), dir_checkpoint + 'CP{}.pth'.format(epoch + 1)) print('Checkpoint {} saved !'.format(epoch + 1))
def train_net(net, epochs=5, batch_size=1, lr=0.1, val_percent=0.5, save_cp=True, gpu=True, img_scale=0.5): dir_img = '/home/buiduchanh/WorkSpace/Unet/Pytorch-UNet/data/data_rust/train_few_images/' dir_mask = '/home/buiduchanh/WorkSpace/Unet/Pytorch-UNet/data/data_rust/train_masks_few_images/' # dir_img = '/home/buiduchanh/WorkSpace/Unet/Pytorch-UNet/data/data_rust/train_images/' # dir_mask = '/home/buiduchanh/WorkSpace/Unet/Pytorch-UNet/data/data_rust/train_masks/' dir_checkpoint = 'checkpoints/' ids = get_ids(dir_img) ids = split_ids(ids) iddataset = split_train_val(ids, val_percent) print(''' Starting training: Epochs: {} Batch size: {} Learning rate: {} Training size: {} Validation size: {} Checkpoints: {} CUDA: {} '''.format(epochs, batch_size, lr, len(iddataset['train']), len(iddataset['val']), str(save_cp), str(gpu))) N_train = len(iddataset['train']) #optimizer = optim.SGD(net.parameters(), # lr=lr, # momentum=0.9, # weight_decay=0.0005) optimizer = optim.Adam(net.parameters(), lr = lr, betas=(0.9, 0.999), eps=1e-08, weight_decay=0, amsgrad=False) criterion = nn.BCELoss() for epoch in range(epochs): print('Starting epoch {}/{}.'.format(epoch + 1, epochs)) net.train() # reset the generators train = get_imgs_and_masks(iddataset['train'], dir_img, dir_mask, img_scale) val = get_imgs_and_masks(iddataset['val'], dir_img, dir_mask, img_scale) epoch_loss = 0 for i, b in enumerate(batch(train, batch_size)): imgs = np.array([i[0] for i in b]).astype(np.float32) true_masks = np.array([i[1] for i in b]) imgs = torch.from_numpy(imgs) true_masks = torch.from_numpy(true_masks) if gpu: imgs = imgs.cuda() true_masks = true_masks.cuda() masks_pred = net(imgs) masks_probs_flat = masks_pred.view(-1) true_masks_flat = true_masks.view(-1) loss = criterion(masks_probs_flat, true_masks_flat) epoch_loss += loss.item() print('{0:.4f} --- loss: {1:.6f}'.format(i * batch_size / N_train, loss.item())) optimizer.zero_grad() loss.backward() optimizer.step() print('Epoch {} finished ! Loss: {}'.format(epoch, epoch_loss / i)) if 1: val_dice = eval_net(net, val, gpu) print('Validation Dice Coeff: {}'.format(val_dice)) if save_cp: torch.save(net.state_dict(), dir_checkpoint + 'CP{}.pth'.format(epoch + 1)) print('Checkpoint {} saved !'.format(epoch + 1))
def train_net(net, epochs=30, batch_size=6, lr=0.1, val_percent=0.05, save_cp=True, gpu=False, img_scale=0.5): # dir_img = 'E:/A_paper_thesis/paper5/tensorflow_deeplabv3plus_scrapingData/dataset/Scraping_Data2/train_db' dir_img = 'data/train_db/' dir_mask = 'data/GT_bw/' dir_checkpoint = 'checkpoint0919/' ids = get_ids(dir_img) ids = split_ids(ids) iddataset = split_train_val(ids, val_percent) print(''' Starting training: Epochs: {} Batch size: {} Learning rate: {} Training size: {} Validation size: {} Checkpoints: {} CUDA: {} '''.format(epochs, batch_size, lr, len(iddataset['train']), len(iddataset['val']), str(save_cp), str(gpu))) N_train = len(iddataset['train']) optimizer = optim.SGD(net.parameters(), lr=lr, momentum=0.9, weight_decay=0.0005) criterion = nn.BCELoss() for epoch in range(epochs): print('Starting epoch {}/{}.'.format(epoch + 1, epochs)) net.train() # reset the generators train = get_imgs_and_masks(iddataset['train'], dir_img, dir_mask, img_scale) val = get_imgs_and_masks(iddataset['val'], dir_img, dir_mask, img_scale) epoch_loss = 0 epoch_iou = 0 epoch_xor=0 for i, b in enumerate(batch(train, batch_size)): imgs = np.array([i[0] for i in b]).astype(np.float32) true_masks = np.array([i[1] for i in b]) imgs = torch.from_numpy(imgs) true_masks = torch.from_numpy(true_masks) if gpu: imgs = imgs.cuda() true_masks = true_masks.cuda() masks_pred = net(imgs) masks_probs_flat = masks_pred.view(-1) true_masks_flat = true_masks.view(-1) loss = criterion(masks_probs_flat, true_masks_flat) epoch_loss += loss.item() print('step:', i) # print('Validation Dice Coeff: {0:.4f} --- loss: {1:.6f}'.format(i * batch_size / N_train, loss.item())) print('Validation Dice Coeff: {0:.4f} --- loss: {1:.6f}'.format(i * batch_size / N_train, loss.item())) # # mean iou # intersect = sum(masks_probs_flat*true_masks_flat) # union = sum(masks_probs_flat+true_masks_flat) # iou = (intersect+0.001)/(union-intersect+0.001) # epoch_iou +=iou # mean iou smooth = 1e-6 # we smooth to avoid our devision 0/0 intersect = sum(masks_probs_flat*true_masks_flat) union = sum(masks_probs_flat+true_masks_flat)-intersect iou = (intersect+smooth)/(union+smooth) epoch_iou +=iou # calculate xor # xor quation is: xor = (union(output hợp ground truth) - intersect(output giao ground truth))/ ground truth # xor = (union-intersect)/ground truth xor = (union - intersect)/sum(true_masks_flat) epoch_xor += xor print('mean IoU: {:.4f}'.format(iou)) # print('mean IoU1: {:.4f}'.format(iou1)) print('mean xor: {:.4f}'.format(xor)) # end of mean iou optimizer.zero_grad() loss.backward() optimizer.step() print('Epoch finished ! epoch_Loss: {:.6f}'.format(epoch_loss / i)) print('epoch_iou: {:.4f}'.format(epoch_iou / i)) print('epoch_xor: {:.4f}'.format(epoch_xor / i)) if 1: val_dice = eval_net(net, val, gpu) print('epoch_Validation Dice Coeff: {:.4f}'.format(val_dice)) # need to write mean iou of evaluate here(reference val_dice) if save_cp: torch.save(net.state_dict(), dir_checkpoint + 'CP{}.pth'.format(epoch + 1)) print('Checkpoint {} saved !'.format(epoch + 1))
def train_net(args, net, val_percent=0.05, save_cp=True): dir_img = os.path.join(args.dataset_folder, 'data/train/') dir_mask = os.path.join(args.dataset_folder, 'data/train_masks/') dir_checkpoint = os.path.join(args.dataset_folder, 'checkpoints/') if not os.path.exists(dir_checkpoint): os.makedirs(dir_checkpoint) ids = get_ids(dir_img) ids = split_ids(ids) iddataset = split_train_val(ids, val_percent) print(''' Starting training: Epochs: {} Batch size: {} Learning rate: {} Training size: {} Validation size: {} Checkpoints: {} '''.format(args.epochs, args.batch_size, args.lr, len(iddataset['train']), len(iddataset['val']), str(save_cp))) N_train = len(iddataset['train']) optimizer = optim.SGD(net.parameters(), lr=args.lr, momentum=0.9, weight_decay=0.0005) criterion = nn.BCELoss() for epoch in range(args.epochs): print('Starting epoch {}/{}.'.format(args.epochs + 1, args.epochs)) net.train() # reset the generators train = get_imgs_and_masks(iddataset['train'], dir_img, dir_mask, args.img_scale) val = get_imgs_and_masks(iddataset['val'], dir_img, dir_mask, args.img_scale) epoch_loss = 0 for i, b in enumerate(batch(train, args.batch_size)): imgs = np.array([i[0] for i in b]).astype(np.float32) true_masks = np.array([i[1] for i in b]) imgs = torch.from_numpy(imgs) true_masks = torch.from_numpy(true_masks) # if gpu: imgs = imgs.cuda() true_masks = true_masks.cuda() masks_pred = net(imgs) masks_probs_flat = masks_pred.view(-1) true_masks_flat = true_masks.view(-1) loss = criterion(masks_probs_flat, true_masks_flat) epoch_loss += loss.item() print('{0:.4f} --- loss: {1:.6f}'.format( i * args.batch_size / N_train, loss.item())) optimizer.zero_grad() loss.backward() optimizer.step() print('Epoch finished ! Loss: {}'.format(epoch_loss / i)) if 1: val_dice = eval_net(net, val) print('Validation Dice Coeff: {}'.format(val_dice)) if save_cp: torch.save(net.state_dict(), dir_checkpoint + 'CP{}.pth'.format(epoch + 1)) print('Checkpoint {} saved !'.format(epoch + 1))
def train_net(net, train_dir=None, groundtruth_dir=None, weight_dir=None, weight_name='DeepInsthink', val_percent=0.05, epochs=5, batch_size=1, lr=0.1, save_cp=True, gpu=False, img_scale=0.5): dir_img = train_dir dir_mask = groundtruth_dir dir_checkpoint = weight_dir ids = get_ids(dir_img) ids = split_ids(ids) iddataset = split_train_val(ids, val_percent) print(''' Starting training: Epochs: {} Batch size: {} Learning rate: {} Training size: {} Validation size: {} Checkpoints: {} CUDA: {} '''.format(epochs, batch_size, lr, len(iddataset['train']), len(iddataset['val']), str(save_cp), str(gpu))) N_train = len(iddataset['train']) optimizer = optim.SGD(net.parameters(), lr=lr, momentum=0.9, weight_decay=0.0005) criterion = nn.BCELoss() max_DSC = 0 max_ep_checkpoint = 0 for epoch in range(epochs): print('Starting epoch {}/{}.'.format(epoch + 1, epochs)) net.train() # reset the generators train = get_imgs_and_masks(iddataset['train'], dir_img, dir_mask, img_scale) val = get_imgs_and_masks(iddataset['val'], dir_img, dir_mask, img_scale) epoch_loss = 0 batchN = 0 for i, b in enumerate(batch(train, batch_size)): imgs = np.array([i[0] for i in b]).astype(np.float32) true_masks = np.array([i[1] for i in b]) imgs = torch.from_numpy(imgs) true_masks = torch.from_numpy(true_masks) if gpu: imgs = imgs.cuda() true_masks = true_masks.cuda() masks_pred = net(imgs) masks_probs_flat = masks_pred.view(-1) true_masks_flat = true_masks.view(-1) loss = criterion(masks_probs_flat, true_masks_flat) epoch_loss += loss.item() optimizer.zero_grad() loss.backward() optimizer.step() print('ep: {3:.0f} batch [{0:.0f}/{1:.0f}] - loss: {2:.6f}'.format( i + 1, N_train / batch_size, loss.item(), epoch + 1)) #if (i % 5==0): # val_dice = eval_net(net, val, gpu) # print('Validation Dice Coeff: {}'.format(val_dice)) val_dice = eval_net(net, val, gpu) print('Epoch {0:} -- Loss: {1:} -- Validation DSC: {2:}'.format( epoch, epoch_loss / i, val_dice)) if (val_dice >= max_DSC): max_DSC = val_dice max_ep_checkpoint = epoch + 1 if save_cp: torch.save( net.state_dict(), dir_checkpoint + weight_name + '-{}.pth'.format(epoch + 1)) print('Checkpoint {} saved !'.format(epoch + 1)) print('Maximum checkpoint is ' + weight_name + '-{0:}' + 'with {1:} DSC'.format(max_ep_checkpoint, max_DSC))
def train_net(net, epochs=20, batch_size=1, lr=0.1, lrd=0.99, val_percent=0.05, save_cp=True, gpu=True, img_scale=0.5, imagepath='', maskpath='', cpsavepath=''): dir_img = imagepath dir_mask = maskpath dir_checkpoint = cpsavepath classweight = [1, 2, 3, 2] ids = get_ids(dir_img) ids = split_ids(ids) iddataset = split_train_val(ids, val_percent) logname = cpsavepath + '/' + 'losslog.txt' print(''' Starting training: Epochs: {} Batch size: {} Learning rate: {} Training size: {} Validation size: {} Checkpoints: {} CUDA: {} '''.format(epochs, batch_size, lr, len(iddataset['train']), len(iddataset['val']), str(save_cp), str(gpu))) N_train = len(iddataset['train']) optimizer = optim.SGD(net.parameters(), lr=lr, momentum=0.9, weight_decay=0.0005) # classweight = [1,4,8,4] criterion = BCELoss_weight(classweight) for epoch in range(epochs): print('Starting epoch {}/{}.'.format(epoch + 1, epochs)) with open(logname, "a") as f: f.write('Starting epoch {}/{}.'.format(epoch + 1, epochs) + "\n") net.train() # reset the generators train = get_imgs_and_masks(iddataset['train'], dir_img, dir_mask, img_scale) val = get_imgs_and_masks(iddataset['val'], dir_img, dir_mask, img_scale) epoch_loss = 0 lr = lr * lrd for param_group in optimizer.param_groups: param_group['lr'] = lr print('lr', lr) for i, b in enumerate(batch(train, batch_size)): imgs = np.array([i[0] for i in b]).astype(np.float32) true_masks = np.array([i[1] for i in b]) true_masks = np.transpose(true_masks, axes=[0, 3, 1, 2]) imgs = torch.from_numpy(imgs) true_masks = torch.from_numpy(true_masks) if gpu: imgs = imgs.cuda() true_masks = true_masks.cuda() masks_pred = net(imgs) # print('masks_pred.shape',masks_pred.shape) # print('true_masks.shape', true_masks.shape) masks_probs_flat = masks_pred true_masks_flat = true_masks loss = criterion(masks_probs_flat, true_masks_flat) epoch_loss += loss.item() printinfo = '{0:.4f} --- loss: {1:.6f}'.format( i * batch_size / N_train, loss.item()) print(printinfo) with open(logname, "a") as f: f.write(printinfo + "\n") optimizer.zero_grad() loss.backward() optimizer.step() print('Epoch finished ! Loss: {}'.format(epoch_loss / i)) with open(logname, "a") as f: f.write('Epoch finished ! Loss: {}'.format(epoch_loss / i) + "\n") if 1: val_dice = eval_net(net, val) print('Validation Dice Coeff: {}'.format(val_dice)) with open(logname, "a") as f: f.write('Validation Dice Coeff: {}'.format(val_dice) + "\n") if save_cp: torch.save(net.state_dict(), dir_checkpoint + 'CP{}.pth'.format(epoch + 1)) print('Checkpoint {} saved !'.format(epoch + 1)) with open(logname, "a") as f: f.write('Checkpoint {} saved !'.format(epoch + 1) + "\n")
def train_net(net, epochs=5, batch_size=1, lr=0.1, val_percent=0.05, save_cp=True, gpu=False, img_scale=0.5): dir_img = 'data/train/' # 训练图像文件夹 dir_mask = 'data/train_masks/' # 图像的结果文件夹 dir_checkpoint = 'checkpoints/' # 训练好的网络保存文件夹 ids = get_ids(dir_img) # 图片名字的后4位为数字,能作为图片id # 得到元祖列表为[(id1,0),(id1,1),(id2,0),(id2,1),...,(idn,0),(idn,1)] # 这样的作用是后面重新设置生成器时会通过后面的0,1作为utils.py中get_square函数的pos参数,pos=0的取左边的部分,pos=1的取右边的部分 # 这样图片的数量就会变成2倍 ids = split_ids(ids) iddataset = split_train_val(ids, val_percent) # 将数据分为训练集和验证集两份 print(''' Starting training: Epochs: {} Batch size: {} Learning rate: {} Training size: {} Validation size: {} Checkpoints: {} CUDA: {} '''.format(epochs, batch_size, lr, len(iddataset['train']), len(iddataset['val']), str(save_cp), str(gpu))) N_train = len(iddataset['train']) # 训练集长度 optimizer = optim.SGD( net.parameters(), # 定义优化器 lr=lr, momentum=0.9, weight_decay=0.0005) criterion = nn.BCELoss() # 损失函数 for epoch in range(epochs): # 开始训练 print('Starting epoch {}/{}.'.format(epoch + 1, epochs)) net.train() # 设置为训练模式 # reset the generators重新设置生成器 # 对输入图片dir_img和结果图片dir_mask进行相同的图片处理,即缩小、裁剪、转置、归一化后,将两个结合在一起,返回(imgs_normalized, masks) train = get_imgs_and_masks(iddataset['train'], dir_img, dir_mask, img_scale) val = get_imgs_and_masks(iddataset['val'], dir_img, dir_mask, img_scale) epoch_loss = 0 for i, b in enumerate(batch(train, batch_size)): imgs = np.array([i[0] for i in b]).astype(np.float32) # 得到输入图像数据 true_masks = np.array([i[1] for i in b]) # 得到图像结果数据 imgs = torch.from_numpy(imgs) true_masks = torch.from_numpy(true_masks) if gpu: imgs = imgs.cuda() true_masks = true_masks.cuda() masks_pred = net(imgs) # 图像输入的网络后得到结果masks_pred,结果为灰度图像 masks_probs_flat = masks_pred.view(-1) # 将结果压扁 true_masks_flat = true_masks.view(-1) loss = criterion(masks_probs_flat, true_masks_flat) # 对两个结果计算损失 epoch_loss += loss.item() print('{0:.4f} --- loss: {1:.6f}'.format(i * batch_size / N_train, loss.item())) optimizer.zero_grad() loss.backward() optimizer.step() print('Epoch finished ! Loss: {}'.format(epoch_loss / i)) # 一次迭代后得到的平均损失 if 1: val_dice = eval_net(net, val, gpu) print('Validation Dice Coeff: {}'.format(val_dice)) if save_cp: torch.save(net.state_dict(), dir_checkpoint + 'CP{}.pth'.format(epoch + 1)) print('Checkpoint {} saved !'.format(epoch + 1))
help='whether to save checkpoint') (options, args) = parser.parse_args() return options if __name__ == '__main__': args = get_args() # dir_img = '/home/zzh/数据/mid project/raw_data' # dir_mask = '/home/zzh/数据/mid project/groundtruth' dir_img = '/home/zhuzhu/Desktop/mid project/raw_data' dir_mask = '/home/zhuzhu/Desktop/mid project/groundtruth' ids = get_ids(dir_img) # 1,2,3,...的生成器 iddataset = split_train_val( ids, args.val_percent) # {'train':[23,98,59,...],'val':[12,37,48,...]} net = UNet(n_channels=1, n_classes=args.num_classes) optimizer = torch.optim.SGD(net.parameters(), lr=args.lr, momentum=0.99, weight_decay=5e-3) scheduler = StepLR(optimizer, step_size=30, gamma=0.1) criterion = nn.BCEWithLogitsLoss() if args.load: print('load model from checkpoint') net.load_state_dict(torch.load('checkpoint/unet.pth')) if args.gpu: net.to(args.device)
if __name__ == '__main__': import numpy as np from utils import split_train_val, batch, get_imgs_and_masks, inputImageViz # get all values that you require from your input args = get_args() dir_img = args.images dir_mask = args.masks dir_checkpoint = 'checkpoints/' val_percent = args.valpercent scale = args.scale maxWidth = args.maxWidth maxHeight = args.maxHeight n_classes= args.classes iddataset = split_train_val(dir_img, val_percent) train = get_imgs_and_masks(iddataset['train'], dir_img, dir_mask, scale, maxWidth, maxHeight, n_classes) val = get_imgs_and_masks(iddataset['val'], dir_img, dir_mask, scale, maxWidth, maxHeight, n_classes) for epoch in range(args.epochs): print('Starting epoch {}/{}.'.format(epoch + 1, epochs)) #net.train() for i, k in enumerate(batch(val, args.batchsize)): print(i) imgs = np.array([i[0] for i in k]).astype(np.float32) true_masks = np.array([i[1] for i in k]) # comment this if required inputImageViz (imgs, true_masks) #------------ your training code here!!!--------------------
def train_net( net, epochs=5, batch_size=1, lr=0.1, val_percent=0.05, # 训练集:验证集= 0.95: 0.05 save_cp=True, gpu=False, img_scale=0.5): dir_img = opt_train.dir_img dir_mask = opt_train.dir_mask dir_checkpoint = opt_train.dir_checkpoint # 得到 图片路径列表 ids为 图片名称(无后缀名) ids = get_ids(dir_img) # 得到truple元组 (无后缀名的 图片名称,序号) # eg:当n为2 图片名称为bobo.jpg 时, 得到(bobo,0) (bobo,1) # 当序号为0 时,裁剪宽度,得到左边部分图片 当序号为1 时,裁剪宽度,得到右边部分图片 ids = split_ids(ids) # 打乱数据集后,按照val_percent的比例来 切分 训练集 和 验证集 iddataset = split_train_val(ids, val_percent) print(''' 开始训练: Epochs: {} Batch size: {} Learning rate: {} 训练集大小: {} 验证集大小: {} GPU: {} '''.format(epochs, batch_size, lr, len(iddataset['train']), len(iddataset['val']), str(gpu))) #训练集大小 N_train = len(iddataset['train']) optimizer = optim.SGD(net.parameters(), lr=lr, momentum=0.9, weight_decay=0.0005) #二进制交叉熵 criterion = nn.BCELoss() for epoch in range(epochs): print('Starting epoch {}/{}.'.format(epoch + 1, epochs)) # reset the generators # 每轮epoch得到 训练集 和 验证集 train = get_imgs_and_masks(iddataset['train'], dir_img, dir_mask, img_scale) val = get_imgs_and_masks(iddataset['val'], dir_img, dir_mask, img_scale) # 重置epoch损失计数器 epoch_loss = 0 for i, b in enumerate(batch(train, batch_size)): # 得到 一个batch的 imgs tensor 及 对应真实mask值 # 当序号为0 时,裁剪宽度,得到左边部分图片[384,384,3] 当序号为1 时,裁剪宽度,得到右边部分图片[384,190,3] imgs = np.array([i[0] for i in b]).astype(np.float32) true_masks = np.array([i[1] for i in b]) # 将值转为 torch tensor imgs = torch.from_numpy(imgs) true_masks = torch.from_numpy(true_masks) # 训练数据转到GPU上 if gpu: imgs = imgs.cuda() true_masks = true_masks.cuda() # 得到 网络输出的预测mask [10,1,384,384] masks_pred = net(imgs) # 经过sigmoid masks_probs = F.sigmoid(masks_pred) masks_probs_flat = masks_probs.view(-1) true_masks_flat = true_masks.view(-1) # 计算二进制交叉熵损失 loss = criterion(masks_probs_flat, true_masks_flat) # 统计一个epoch的所有batch的loss之和,用以计算 一个epoch的 loss均值 epoch_loss += loss.item() # 输出 当前epoch的第几个batch 及 当前batch的loss print('{0:.4f} --- loss: {1:.6f}'.format(i * batch_size / N_train, loss.item())) # 优化器梯度清零 optimizer.zero_grad() # 反向传播 loss.backward() # 更新参数 optimizer.step() # 一轮epoch结束,该轮epoch的 loss均值 print('Epoch finished ! Loss: {}'.format(epoch_loss / i)) # 每轮epoch之后使用验证集进行评价 if True: # 评价函数:Dice系数 Dice距离用于度量两个集合的相似性 val_dice = eval_net(net, val, gpu) print('Validation Dice Coeff: {}'.format(val_dice)) # 保存模型 if save_cp: torch.save(net.state_dict(), dir_checkpoint + 'CP{}.pth'.format(epoch + 1)) print('Checkpoint {} saved !'.format(epoch + 1))
def train_net(net, epochs=5, batch_size=1, lr=0.1, val_percent=0.05, save_cp=True, gpu=False, img_scale=0.5): dir_img = '../dataset/train/images/' dir_mask = '../dataset/train/masks/' dir_checkpoint = 'checkpoints/' ids = get_ids(dir_img) ids = split_ids(ids) iddataset = split_train_val(ids, val_percent) print((''' Starting training: Epochs: {} Batch size: {} Learning rate: {} Training size: {} Validation size: {} Checkpoints: {} CUDA: {} '''.format(epochs, batch_size, lr, len(iddataset['train']), len(iddataset['val']), str(save_cp), str(gpu)))) N_train = len(iddataset['train']) optimizer = optim.SGD(net.parameters(), lr=lr, momentum=0.9, weight_decay=0.0005) criterion = nn.BCELoss() for epoch in range(epochs): print(('Starting epoch {}/{}.'.format(epoch + 1, epochs))) net.train() # reset the generators train = get_imgs_and_masks(iddataset['train'], dir_img, dir_mask, img_scale) val = get_imgs_and_masks(iddataset['val'], dir_img, dir_mask, img_scale) epoch_loss = 0 for i, b in enumerate(batch(train, batch_size)): # b[batch_id][0/1]: a batch of image(0)+mask(1) #print(('b[0]',b[0][0].shape,b[0][1].shape)) #imgs = [] #for img_msk in b: # imgs.append(img_msk[0]) #print(len(imgs)) #imgs = np.array(imgs) # Wrong: not all images are of the same shape imgs = np.array([i[0] for i in b]).astype(np.float32) true_masks = np.array([i[1] for i in b]) imgs = torch.from_numpy(imgs) true_masks = torch.from_numpy(true_masks) if gpu: imgs = imgs.cuda() true_masks = true_masks.cuda() masks_pred = net(imgs) masks_probs_flat = masks_pred.view(-1) true_masks_flat = true_masks.view(-1) loss = criterion(masks_probs_flat, true_masks_flat) epoch_loss += loss.item() print(('{0:.4f} --- loss: {1:.6f}'.format(i * batch_size / N_train, loss.item()))) optimizer.zero_grad() loss.backward() optimizer.step() print(('Epoch finished ! Loss: {}'.format(epoch_loss / i))) if 1: val_dice = eval_net(net, val, gpu) print(('Validation Dice Coeff: {}'.format(val_dice))) if save_cp: torch.save(net.state_dict(), dir_checkpoint + 'cropped_CP{}.pth'.format(epoch + 1)) print(('Checkpoint {} saved !'.format(epoch + 1)))
def train_net(net, epochs=5, batch_size=1, lr=0.1, val_percent=0.2, save_cp=True, gpu=False, img_scale=0.5): path = [['data/ori1/', 'data/gt1/'], ['data/original1/', 'data/ground_truth1/'], ['data/Original/', 'data/Ground_Truth/']] dir_img = path[0][0] dir_mask = path[0][1] dir_checkpoint = 'sdgcheck/' ids = get_ids(dir_img) ids = split_ids(ids) iddataset = split_train_val(ids, val_percent) print(''' Starting training: Epochs: {} Batch size: {} Learning rate: {} Training size: {} Validation size: {} Checkpoints: {} CUDA: {} '''.format(epochs, batch_size, lr, len(iddataset['train']), len(iddataset['val']), str(save_cp), str(gpu))) N_train = len(iddataset['train']) optimizer = optim.SGD(net.parameters(), lr=lr, momentum=0.7, weight_decay=0.005) ''' optimizer = optim.Adam(net.parameters(), lr=lr, weight_decay=0.0005) ''' criterion = nn.BCELoss() for epoch in range(epochs): print('Starting epoch {}/{}.'.format(epoch + 1, epochs)) net.train() # reset the generators train = get_imgs_and_masks(iddataset['train'], dir_img, dir_mask, img_scale) val = get_imgs_and_masks(iddataset['val'], dir_img, dir_mask, img_scale) epoch_loss = 0 x = 0 for i, b in enumerate(batch(train, batch_size)): imgs = np.array([i[0] for i in b]).astype(np.float32) true_masks = np.array([i[1] for i in b]) ''' ori=np.transpose(imgs[0], axes=[1, 2, 0]) scipy.misc.imsave("ori/ori_"+str(x)+'.jpg', ori) gt = np.stack((true_masks[0],)*3, axis=-1) #gt=np.transpose(true_masks[0], axes=[1, 2, 0]) scipy.misc.imsave("gt/gt_"+str(x)+'.jpg', gt) ''' imgs = torch.from_numpy(imgs) true_masks = torch.from_numpy(true_masks) x += 1 if gpu: imgs = imgs.cuda() true_masks = true_masks.cuda() masks_pred = net(imgs) masks_probs_flat = masks_pred.view(-1) true_masks_flat = true_masks.view(-1) loss = criterion(masks_probs_flat, true_masks_flat) epoch_loss += loss.item() print('{0:.4f} --- loss: {1:.6f}'.format(i * batch_size / N_train, loss.item())) optimizer.zero_grad() loss.backward() optimizer.step() print('Epoch finished ! Loss: {}'.format(epoch_loss / i)) if 1: val_dice = eval_net(net, val, gpu) print('Validation Dice Coeff: {}'.format(val_dice)) if save_cp: torch.save(net.state_dict(), dir_checkpoint + 'CP{}.pth'.format(epoch + 1)) print('Checkpoint {} saved !'.format(epoch + 1))
def train_net(net, device, epochs=5, batch_size=1, lr=0.1, val_percent=0.15, save_cp=True, img_scale=0.5): ids = get_ids(dir_img) iddataset = split_train_val(ids, val_percent) n_train = len(iddataset['train']) n_val = len(iddataset['val']) optimizer = optim.Adam(net.parameters(), lr=lr) if net.n_classes > 1: criterion = nn.CrossEntropyLoss() else: criterion = nn.BCEWithLogitsLoss() for epoch in range(epochs): net.train() # reset the generators train = get_imgs_and_masks(iddataset['train'], dir_img, dir_mask, img_scale) val = get_imgs_and_masks(iddataset['val'], dir_img, dir_mask, img_scale) epoch_loss = 0 with tqdm(total=n_train, desc='Epoch {0}/{1}'.format(epoch + 1, epochs), unit='img') as pbar: for i, b in enumerate(batch(train, batch_size)): current_lr = adjust_learning_rate(optimizer, epoch, epochs, pbar.n, n_train) random_rate = 0 if epoch > epochs / 2: random_rate = (epoch * 0.1) / epochs b = custom_transforms.random_data_augmentation( b, random_rate=random_rate) imgs = np.array([i[0] for i in b]).astype(np.float32) true_masks = np.array([i[1][0] for i in b]) imgs = torch.from_numpy(imgs) true_masks = torch.from_numpy(true_masks) imgs = imgs.to(device=device) true_masks = true_masks.to(device=device) masks_pred = net(imgs) loss = criterion(masks_pred, true_masks.long()) epoch_loss += loss.item() pbar.set_postfix( **{ 'lr:{0}, random_rate:{1}, loss:'.format( current_lr, random_rate): loss.item() }) optimizer.zero_grad() loss.backward() optimizer.step() pbar.update(batch_size) if save_cp: try: os.mkdir(dir_checkpoint) logging.info('Created checkpoint directory') except OSError: pass torch.save(net.state_dict(), dir_checkpoint + 'CP_epoch{0}.pth'.format(epoch + 1)) logging.info('Checkpoint {0} saved !'.format(epoch + 1)) val_score = eval_net(net, val, device, n_val) if net.n_classes > 1: logging.info('Validation cross entropy: {0}'.format(val_score)) else: logging.info('Validation Dice Coeff: {0}'.format(val_score))
"gpu": args.gpu, "load": args.load, "channel_txt": args.channel_txt, "scale": args.scale, "lr": args.lr, "iters": args.iters, "epochs": args.epochs }, indent=4, sort_keys=True))) # Dataset if not os.path.exists(splitfile): # Our constant datasplit ids = get_ids(dir_img) # [file1, file2] ids = split_ids(ids) # [(file1, 0), (file1, 1), (file2, 0), ...] iddataset = split_train_val(ids, 0.2, splitfile) log.info("New split dataset") else: with open(splitfile) as f: iddataset = json.load(f) log.info("Load split dataset") train = get_imgs_and_masks(iddataset['train'], dir_img, dir_mask, args.scale) val = get_imgs_and_masks(iddataset['val'], dir_img, dir_mask, args.scale) # Model Initialization net = UNet(n_channels=3, n_classes=1, f_channels=args.channel_txt) log.info("Built model using {}...".format(args.channel_txt)) if args.gpu:
def train_net(net, device, epochs=5, batch_size=1, lr=0.1, val_percent=0.15, save_cp=True, img_scale=0.5): ids = get_ids(dir_img) iddataset = split_train_val(ids, val_percent) logging.info('''Starting training: Epochs: {epochs} Batch size: {batch_size} Learning rate: {lr} Training size: {len(iddataset["train"])} Validation size: {len(iddataset["val"])} Checkpoints: {save_cp} Device: {device.type} Images scaling: {img_scale} ''') n_train = len(iddataset['train']) n_val = len(iddataset['val']) optimizer = optim.Adam(net.parameters(), lr=lr) if net.n_classes > 1: criterion = nn.CrossEntropyLoss() else: criterion = nn.BCEWithLogitsLoss() for epoch in range(epochs): net.train() # reset the generators train = get_imgs_and_masks(iddataset['train'], dir_img, dir_mask, img_scale) val = get_imgs_and_masks(iddataset['val'], dir_img, dir_mask, img_scale) epoch_loss = 0 f1_score = 0 num = 0 with tqdm(total=n_train, desc='Epoch {epoch + 1}/{epochs}', unit='img') as pbar: for i, b in enumerate(batch(train, batch_size)): imgs = np.array([i[0] for i in b]).astype(np.float32) true_masks = np.array([i[1] for i in b]) imgs = torch.from_numpy(imgs) true_masks = torch.from_numpy(true_masks) imgs = imgs.to(device=device) true_masks = true_masks.to(device=device) masks_pred = net(imgs) # print('mask:',masks_pred.size()) # print('lab:',true_masks.size()) loss = criterion(masks_pred, true_masks) masks_pred_np = masks_pred.detach().cpu().numpy() true_masks_np = true_masks.detach().cpu().numpy() epoch_loss += loss.item() # print("----------------------------------------") # print('masks_pred',type(masks_pred),masks_pred,'\n') # print('true_masks',type(true_masks),true_masks,'\n') # print('mask:',masks_pred.size(),'\n') # print('lab:',true_masks.size(),'\n') pre_2D = np.array(masks_pred_np[0][0]) true_2D = np.array(true_masks_np[0][0]) pre_2D_threhold = pre_2D pre_2D_threhold[pre_2D_threhold > 0.5] = 1 pre_2D_threhold[pre_2D_threhold <= 0.5] = 0 # print("pre_2D.shape",pre_2D.shape,'\n') # print("true_2D.shape" ,true_2D.shape,'\n') # print("true_2D.flatten()",true_2D.flatten(),'\n') # print("pre_2D.flatten()",pre_2D.flatten(),'\n') pixel_accuracy = (pre_2D, true_2D) f1_score += metrics.f1_score(true_2D.flatten(), pre_2D_threhold.flatten()) num = num + 1 # print("----------------------------------------") # val_score1 = eval_net1(net,val,device,n_val) pbar.set_postfix(**{'loss (batch)': loss.item()}) optimizer.zero_grad() loss.backward() optimizer.step() pbar.update(batch_size) if save_cp: try: os.mkdir(dir_checkpoint) logging.info('Created checkpoint directory') except OSError: pass torch.save(net.state_dict(), dir_checkpoint + 'CP_epoch{epoch + 1}.pth') logging.info('Checkpoint {epoch + 1} saved !') val_score = eval_net(net, val, device, n_val) f1_score /= num print("f1-score:", f1_score, '\n') if net.n_classes > 1: logging.info('Validation cross entropy: {}'.format(val_score)) else: logging.info('Validation Dice Coeff: {}'.format(val_score))
def train_net(net, epochs=5, batch_size=1, lr=0.1, val_percent=0.05, save_cp=True, gpu=False, img_scale=0.5): # dir_img = 'data/train/' # dir_mask = 'data/train_masks/' dir_img = 'E:/git/dataset/tgs-salt-identification-challenge/train/images/' dir_mask = 'E:/git/dataset/tgs-salt-identification-challenge/train/masks/' # dir_img = 'E:/git/dataset/tgs-salt-identification-challenge/train/my_images/' # dir_mask = 'E:/git/dataset/tgs-salt-identification-challenge/train/my_masks/' dir_checkpoint = 'checkpoints/' ids = get_ids(dir_img) ids = split_ids(ids) iddataset = split_train_val(ids, val_percent) print(''' Starting training: Epochs: {} Batch size: {} Learning rate: {} Training size: {} Validation size: {} Checkpoints: {} CUDA: {} '''.format(epochs, batch_size, lr, len(iddataset['train']), len(iddataset['val']), str(save_cp), str(gpu))) N_train = len(iddataset['train']) optimizer = optim.SGD(net.parameters(), lr=lr, momentum=0.9, weight_decay=0.0005) criterion = nn.BCELoss() for epoch in range(epochs): print('Starting epoch {}/{}.'.format(epoch + 1, epochs)) net.train() # reset the generators train = get_imgs_and_masks(iddataset['train'], dir_img, dir_mask, img_scale) val = get_imgs_and_masks(iddataset['val'], dir_img, dir_mask, img_scale) epoch_loss = 0 for i, b in enumerate(batch(train, batch_size)): imgs = np.array([i[0] for i in b]).astype(np.float32) # true_masks = np.array([i[1] for i in b])#np.rot90(m) true_masks = np.array([i[1].T / 65535 for i in b]) #np.rot90(m) # show_batch_image(true_masks) imgs = torch.from_numpy(imgs) true_masks = torch.from_numpy(true_masks) if gpu: imgs = imgs.cuda() true_masks = true_masks.cuda() # show_batch_image(imgs) masks_pred = net(imgs) masks_probs_flat = masks_pred.view(-1) true_masks_flat = true_masks.view(-1) loss = criterion(masks_probs_flat, true_masks_flat) epoch_loss += loss.item() print('{0:.4f} --- loss: {1:.6f}'.format(i * batch_size / N_train, loss.item())) optimizer.zero_grad() loss.backward() optimizer.step() print('Epoch finished ! Loss: {}'.format(epoch_loss / i)) if 1: val_dice = eval_net(net, val, gpu) print('Validation Dice Coeff: {}'.format(val_dice)) if save_cp: torch.save(net.state_dict(), dir_checkpoint + 'CP{}.pth'.format(epoch + 1)) print('Checkpoint {} saved !'.format(epoch + 1))
def get_dataloaders(meter_type, batch_size=20, trial=False): train_records, test_records = get_train_test_records(meter_type) with open(log_file, "a") as f: f.write("{}: train = {}, test = {}\n".format("ABCDE"[meter_type], len(train_records), len(test_records))) if meter_type == 2: batch_size //= 2 pass train_records, val_records = utils.split_train_val(train_records) if trial: n = 100 trainset = PointerMeters(meter_type, train_records[:n], transform=torchvision.transforms.Compose([ RandomCrop(meter_type), Resize(meter_type), ToTensor() ])) valset = PointerMeters(meter_type, val_records[:n], transform=torchvision.transforms.Compose([ RandomCrop(meter_type), Resize(meter_type), ToTensor() ])) else: trainset = PointerMeters(meter_type, train_records, transform=torchvision.transforms.Compose([ RandomCrop(meter_type), Resize(meter_type), ToTensor() ])) valset = PointerMeters(meter_type, val_records, transform=torchvision.transforms.Compose([ RandomCrop(meter_type), Resize(meter_type), ToTensor() ])) train_loader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, shuffle=True) val_loader = torch.utils.data.DataLoader(valset, batch_size=batch_size, shuffle=False) with open(log_file, "a") as f: f.write("train: {}, val: {}\n".format(len(train_loader.dataset), len(val_loader.dataset))) show_samples(train_loader, "train") if trial: np.random.shuffle(test_records) testset = PointerMeters(meter_type, test_records[:n], transform=torchvision.transforms.Compose([ CentralCrop(meter_type), Resize(meter_type), ToTensor() ])) else: testset = PointerMeters(meter_type, test_records, transform=torchvision.transforms.Compose([ CentralCrop(meter_type), Resize(meter_type), ToTensor() ])) test_loader = torch.utils.data.DataLoader(testset, batch_size=batch_size, shuffle=True) with open(log_file, "a") as f: f.write("test: {}\n".format(len(test_loader.dataset))) show_samples(test_loader, "test") return train_loader, val_loader, test_loader