partition = 'train'
shape_train = ShapeData(ROOT_DIR, partition)
train_loader = torch.utils.data.DataLoader(shape_train,
batch_size=10,
shuffle=True,
)
partition = 'val'
shape_val = ShapeData(ROOT_DIR, partition)
val_loader = torch.utils.data.DataLoader(shape_val,
batch_size=10,
shuffle=False
)
# Create model
model = ShapeUNet((1, 200,200))
model.to(device)
# Specify optimizer and criterion
lr = 1e-4
optimizer = Adam(model.parameters(), lr=lr)
NUM_OF_EPOCHS = 50
#training
train(model, train_loader, val_loader, optimizer, NUM_OF_EPOCHS, 'weights/')
# PATH_TO_SAVE = 'predictions/'
# preds = []
# masks = []
# imgs = []
# Create Data Loaders
partition = 'train'
ultrasound_train = UltrasoundData(ROOT_DIR, partition)
train_loader = torch.utils.data.DataLoader(
ultrasound_train,
batch_size=1,
shuffle=True,
)
partition = 'val'
ultrasound_val = UltrasoundData(ROOT_DIR, partition)
val_loader = torch.utils.data.DataLoader(ultrasound_val,
batch_size=1,
shuffle=False)
# # Create models
PATH_TO_SHAPE_WEIGHT_MODEL = 'weights_shape_net/Unet_manually_corrupted_lumen0.113253.pth'
# Specify optimizer and criterion
NUM_OF_EPOCHS = 50
unet = UNet((3, 544, 544))
shapeUNet = ShapeUNet((4, 544, 544))
shapeUNet.load_state_dict(torch.load(PATH_TO_SHAPE_WEIGHT_MODEL))
unet.to(device)
shapeUNet.to(device)
lr = 1e-4
optimizer1 = Adam(shapeUNet.parameters(), lr=lr)
optimizer2 = Adam(unet.parameters(), lr=lr)
train(unet, shapeUNet, train_loader, val_loader, optimizer1, optimizer2,
NUM_OF_EPOCHS, 'weights_shape_net/')