#np.random.shuffle(indices) print('There is no random shuffle: initial portion of the dataset is used for train and the last portion for validation') train_indices, test_indices = indices[split:], indices[:split] train_sampler = SubsetRandomSampler(train_indices) test_sampler = SubsetRandomSampler(test_indices) train_loader = DataLoader(dataset, batch_size=batch_size_train, sampler=train_sampler) test_loader = DataLoader(dataset, batch_size=batch_size_val, sampler=test_sampler) print("Number of training/test patches:", (len(train_indices),len(test_indices))) # network net = ProbabilisticUnet(input_channels=1, num_classes=1, num_filters=[32,64,128,192], latent_dim=2, no_convs_fcomb=4, beta=10.0) net.cuda() # optimizer optimizer = torch.optim.Adam(net.parameters(), lr=lr, weight_decay=l2_reg) secheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=lr_decay_every, gamma=lr_decay) # logging train_loss = [] test_loss = [] best_val_loss = 999.0 for epoch in range(epochs): net.train() loss_train = 0 loss_segmentation = 0 # training loop for step, (patch, mask, _) in enumerate(train_loader): patch = patch.cuda() mask = mask.cuda()
train_loader = DataLoader(dataset, batch_size=train_batch_size, sampler=train_sampler) # 训练 train_eval_loader = DataLoader(dataset, batch_size=test_batch_size, sampler=train_sampler) # 评估 test_loader = DataLoader(dataset, batch_size=test_batch_size, sampler=test_sampler) # 评估 print("Number of training/test patches:", (len(train_indices),len(test_indices))) # 网络模型 net = ProbabilisticUnet(input_channels=1, num_classes=class_num, num_filters=[32,64,128,192], latent_dim=2, no_convs_fcomb=4, beta=10.0) net.to(device) # 优化器 optimizer = torch.optim.Adam(net.parameters(), lr=learning_rate, weight_decay=0) # 训练模型并保存 try: # 训练 for epoch in range(epochs): print("Epoch {}".format(epoch)) # 训练 net.train() losses = 0 # 计算平均loss值 for step, (patch, mask, _) in enumerate(train_loader): patch = patch.to(device) mask = mask.to(device) # mask = torch.unsqueeze(mask,1) (batch_size,240,240)->(batch_size,1,240,240) net.forward(patch, mask, training=True)
np.random.shuffle(indices) train_indices, test_indices = indices[split:], indices[:split] train_sampler = SubsetRandomSampler(train_indices) test_sampler = SubsetRandomSampler(test_indices) train_loader = DataLoader(dataset, batch_size=5, sampler=train_sampler) test_loader = DataLoader(dataset, batch_size=1, sampler=test_sampler) print("Number of training/test patches:", (len(train_indices), len(test_indices))) net = ProbabilisticUnet(input_channels=1, num_classes=1, num_filters=[32, 64, 128, 192], latent_dim=2, no_convs_fcomb=4, beta=10.0) net.to(device) optimizer = torch.optim.Adam(net.parameters(), lr=1e-4, weight_decay=0) epochs = 10 for epoch in range(epochs): for step, (patch, mask, _) in enumerate(train_loader): patch = patch.to(device) mask = mask.to(device) mask = torch.unsqueeze(mask, 1) net.forward(patch, mask, training=True) elbo = net.elbo(mask) reg_loss = l2_regularisation(net.posterior) + l2_regularisation( net.prior) + l2_regularisation(net.fcomb.layers) loss = -elbo + 1e-5 * reg_loss optimizer.zero_grad() loss.backward() optimizer.step()
def train(args): num_epoch = args.epoch learning_rate = args.learning_rate task_dir = args.task trainset = MedicalDataset(task_dir=task_dir, mode='train' ) validset = MedicalDataset(task_dir=task_dir, mode='valid') model = ProbabilisticUnet(input_channels=1, num_classes=1, num_filters=[32,64,128,192], latent_dim=2, no_convs_fcomb=4, beta=10.0) model.to(device) #summary(model, (1,320,320)) optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate, weight_decay=0) criterion = torch.nn.BCELoss() for epoch in range(num_epoch): model.train() while trainset.iteration < args.iteration: x, y = trainset.next() x, y = torch.from_numpy(x).unsqueeze(0).cuda(), torch.from_numpy(y).unsqueeze(0).cuda() #print(x.size(), y.size()) #output = torch.nn.Sigmoid()(model(x)) model.forward(x,y,training=True) elbo = model.elbo(y) reg_loss = l2_regularisation(model.posterior) + l2_regularisation(model.prior) + l2_regularisation(model.fcomb.layers) loss = -elbo + 1e-5 * reg_loss #loss = criterion(output, y) optimizer.zero_grad() loss.backward() optimizer.step() trainset.iteration = 0 model.eval() with torch.no_grad(): while validset.iteration < args.test_iteration: x, y = validset.next() x, y = torch.from_numpy(x).unsqueeze(0).cuda(), torch.from_numpy(y).unsqueeze(0).cuda() #output = torch.nn.Sigmoid()(model(x, y)) model.forward(x,y,training=True) elbo = model.elbo(y) reg_loss = l2_regularisation(model.posterior) + l2_regularisation(model.prior) + l2_regularisation(model.fcomb.layers) valid_loss = -elbo + 1e-5 * reg_loss validset.iteration = 0 print('Epoch: {}, elbo: {:.4f}, regloss: {:.4f}, loss: {:.4f}, valid loss: {:.4f}'.format(epoch+1, elbo.item(), reg_loss.item(), loss.item(), valid_loss.item())) """ #Logger # 1. Log scalar values (scalar summary) info = { 'loss': loss.item(), 'accuracy': valid_loss.item() } for tag, value in info.items(): Logger.scalar_summary(tag, value, epoch+1) # 2. Log values and gradients of the parameters (histogram summary) for tag, value in model.named_parameters(): tag = tag.replace('.', '/') Logger.histo_summary(tag, value.data.cpu().numpy(), epoch+1) Logger.histo_summary(tag+'/grad', value.grad.data.cpu().numpy(), epoch+1) """ torch.save(model.state_dict(), './save/'+trainset.task_dir+'model.pth')