def train_model(model, data_train, criterion, optimizer):
"""Train the model and report validation error with training error
Args:
model: the model to be trained
criterion: loss function
data_train (DataLoader): training dataset
"""
model.train()
for batch, (images, masks) in enumerate(data_train):
images = Variable(images.cuda())
masks = Variable(masks.cuda())
outputs = model(images).cuda()
# print(masks.shape, outputs.shape)
loss = criterion(outputs, masks)
optimizer.zero_grad()
loss.backward()
# Update weights
optimizer.step()
# total_loss = get_loss_train(model, data_train, criterion)
def get_loss_train(model, data_train, criterion):
"""
Calculate loss over train set
"""
model.eval()
total_acc = 0
total_loss = 0
total_dice=0
for batch, (images, masks) in enumerate(data_train):
with torch.no_grad():
images = Variable(images.cuda())
masks = Variable(masks.cuda())
outputs = model(images).cuda()
loss = criterion(outputs, masks)
preds = torch.argmax(outputs, dim=1).float()
acc = accuracy_check_for_batch(masks.cpu(), preds.cpu(), images.size()[0])
total_acc = total_acc + acc
total_loss = total_loss + loss.cpu().item()
return total_acc/(batch+1), total_loss/(batch + 1)
def validate_model(model, data_val, criterion, epoch, make_prediction=True, save_folder_name='prediction'):
"""
Validation run
"""
# calculating validation loss
total_val_loss = 0
total_val_acc = 0
for batch, (images_v, masks_v, original_msk) in enumerate(data_val):
stacked_img = torch.Tensor([]).cuda()
for index in range(images_v.size()[1]):
with torch.no_grad():
image_v = Variable(images_v[:, index, :, :].unsqueeze(0).cuda())
mask_v = Variable(masks_v[:, index, :, :].squeeze(1).cuda())
# print(image_v.shape, mask_v.shape)
output_v = model(image_v)
total_val_loss = total_val_loss + criterion(output_v, mask_v).cpu().item()
# print('out', output_v.shape)
output_v = torch.argmax(output_v, dim=1).float()
stacked_img = torch.cat((stacked_img, output_v))
if make_prediction:
im_name = batch # TODO: Change this to real image name so we know
pred_msk = save_prediction_image(stacked_img, im_name, epoch, save_folder_name)
acc_val = accuracy_check(original_msk, pred_msk)
total_val_acc = total_val_acc + acc_val
return total_val_acc/(batch + 1), total_val_loss/((batch + 1)*4)
def test_model(model_path, data_test, epoch, save_folder_name='prediction'):
"""
Test run
"""
model = torch.load(model_path)
model = torch.nn.DataParallel(model, device_ids=list(
range(torch.cuda.device_count()))).cuda()
model.eval()
for batch, (images_t) in enumerate(data_test):
stacked_img = torch.Tensor([]).cuda()
for index in range(images_t.size()[1]):
with torch.no_grad():
image_t = Variable(images_t[:, index, :, :].unsqueeze(0).cuda())
# print(image_v.shape, mask_v.shape)
output_t = model(image_t)
output_t = torch.argmax(output_t, dim=1).float()
stacked_img = torch.cat((stacked_img, output_t))
im_name = batch # TODO: Change this to real image name so we know
_ = save_prediction_image(stacked_img, im_name, epoch, save_folder_name)
print("Finish Prediction!")
def save_prediction_image(stacked_img, im_name, epoch, save_folder_name="result_images", save_im=True):
"""save images to save_path
Args:
stacked_img (numpy): stacked cropped images
save_folder_name (str): saving folder name
"""
div_arr = division_array(388, 2, 2, 512, 512)
img_cont = image_concatenate(stacked_img.cpu().data.numpy(), 2, 2, 512, 512)
img_cont = polarize((img_cont)/div_arr)*255
img_cont_np = img_cont.astype('uint8')
img_cont = Image.fromarray(img_cont_np)
# organize images in every epoch
desired_path = save_folder_name + '/epoch_' + str(epoch) + '/'
# Create the path if it does not exist
if not os.path.exists(desired_path):
os.makedirs(desired_path)
# Save Image!
export_name = str(im_name) + '.png'
img_cont.save(desired_path + export_name)
return img_cont_np
def polarize(img):
''' Polarize the value to zero and one
Args:
img (numpy): numpy array of image to be polarized
return:
img (numpy): numpy array only with zero and one
'''
img[img >= 0.5] = 1
img[img < 0.5] = 0
return img
model = CleanU_Net(in_channels=1, out_channels=2).cuda()
#model = UNet_model(1,1).cuda()
#model.weight_init()
model = torch.nn.DataParallel(model, device_ids=list(range(torch.cuda.device_count()))).cuda()
criterion = nn.CrossEntropyLoss()
# Waits for everything to finish running
torch.cuda.synchronize()
header = ['epoch', 'train loss', 'train acc', 'val loss', 'val acc']
save_file_name = "/content/drive/My Drive/APP/DATA/data/history/RMS prop.csv"
save_dir = "/content/drive/My Drive/APP/DATA/data/history"
model_save_dir = "/content/drive/My Drive/APP/DATA/data/history/savemodel"
image_save_path = "/content/drive/My Drive/APP/DATA/data/history/model imag"
class DiceCoeffLoss(nn.Module):
def __init__(self,smooth=1):
super(DiceCoeffLoss, self).__init__()
self.smooth = smooth
def forward(self,input, target):
iflat = input.view(-1).float()
tflat = target.view(-1).float()
intersection = (iflat * tflat).sum().float()
return 1 - ((2. * intersection + self.smooth) / (iflat.sum() + tflat.sum() + self.smooth))
# Loss function
# Optimizerd
optimizer = torch.optim.Adam(model.parameters(),lr=0.004,weight_decay=1e-3)
# Parameters
epoch_start = 0
epoch_end = 1000
print("Initializing Training!")
for i in range(epoch_start, epoch_end):
# train the model
train_model(model, SEM_train_load, criterion, optimizer)
train_acc, train_loss= get_loss_train(model, SEM_train_load, criterion)
#train_loss = train_loss / len(SEM_train)
print('Epoch', str(i+1), 'Train loss:', train_loss, "Train acc", train_acc )
# Validation every 5 epoch
if (i+1) % 5 == 0:
val_acc, val_loss = validate_model(
model, SEM_val_load, criterion, i+1, True, image_save_path)
print('Val loss:', val_loss, "val acc:", val_acc)
values = [i+1, train_loss, train_acc, val_loss, val_acc]
export_history(header, values, save_dir, save_file_name)
if (i+1) % 100 == 0: # save model every 10 epoch
save_models(model, model_save_dir, i+1)
torch.cuda.synchronize()
in the code above my training network how can i calculate the AUC_ROC ,AUC_PR ,Dice_coeff for the dataset?