# Create data loaders (return data in batches)
trainset_loader, valset_loader = \
data.get_train_val_loaders(train_dir=args.train_dir,
max_trainset_size=args.max_trainset_size,
collate_fn=csv_collator,
height=args.height,
width=args.width,
seed=args.seed,
batch_size=args.batch_size,
drop_last_batch=args.drop_last_batch,
num_workers=args.nThreads,
val_dir=args.val_dir,
max_valset_size=args.max_valset_size)
# Model
with peter('Building network'):
model = unet_model.UNet(3,
1,
height=args.height,
width=args.width,
known_n_points=args.n_points,
device=device)
num_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f" with {ballpark(num_params)} trainable parameters. ", end='')
model = nn.DataParallel(model)
model.to(device)
# Loss functions
loss_regress = nn.SmoothL1Loss()
loss_loc = losses.WeightedHausdorffDistance(resized_height=args.height,
resized_width=args.width,
testset_loader = data.DataLoader(testset,
batch_size=1,
num_workers=args.nThreads,
collate_fn=csv_collator)
# Array with [height, width] of the new size
resized_size = np.array([args.height, args.width])
# Loss function
criterion_training = losses.WeightedHausdorffDistance(resized_height=args.height,
resized_width=args.width,
return_2_terms=True,
device=device)
# Restore saved checkpoint (model weights)
with peter("Loading checkpoint"):
# Pretrained models that come with this package
if args.model == 'unet_256x256_sorghum':
args.model = os.path.join(os.path.dirname(os.path.realpath(__file__)),
'checkpoints',
'unet_256x256_sorghum.ckpt')
if os.path.isfile(args.model):
if args.cuda:
checkpoint = torch.load(args.model)
else:
checkpoint = torch.load(
args.model, map_location=lambda storage, loc: storage)
# Model
if args.n_points is None:
if 'n_points' not in checkpoint:
training_transforms += [RandomVerticalFlipImageAndLabel(p=0.5)]
training_transforms += [ScaleImageAndLabel(size=(args.height, args.width))]
training_transforms += [transforms.ToTensor()]
training_transforms += [transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]
trainset = CSVDataset(args.train_dir,
transforms=transforms.Compose(training_transforms),
max_dataset_size=args.max_trainset_size)
trainset_loader = DataLoader(trainset,
batch_size=args.batch_size,
drop_last=args.drop_last_batch,
shuffle=True,
num_workers=args.nThreads,
collate_fn=csv_collator)
# Model
with peter('Building network'):
model = unet_model.UNet(3,
1,
height=args.height,
width=args.width,
known_n_points=args.n_points)
num_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f" with {ballpark(num_params)} trainable parameters. ", end='')
model = nn.DataParallel(model)
model.to(device)
# Loss function
loss_regress = nn.SmoothL1Loss()
loss_loc = losses.WeightedHausdorffDistance(resized_height=args.height,
resized_width=args.width,
p=args.p,
config = utils.config_parser(config_path, experiment_type="training")
torch.set_default_dtype(torch.float32)
device_cpu = torch.device('cpu')
device = torch.device('cuda:0') if config['use_cuda'] else device_cpu
# data_dictionary,batch_size,num_workers,instance_seg = False):
test_loader = cranberry_dataset.build_single_loader(
data_dictionary=config['data']['eval_dir'],
batch_size=config['testing']['batch_size'],
num_workers=config['testing']['num_workers'],
instance_seg=config['data']['instance_seg'],
test=config['testing']['img_only'],
has_mask=config['data']['has_mask'])
with peter('Building Network'):
model = unet_refined.UNetRefined(n_channels=3, n_classes=2)
# model = unet_regres.Unet(in_channels=3,classes=2,decoder_channels= (512,256,128),encoder_depth=3)
num_params = sum(p.numel() for p in model.parameters()
if p.requires_grad)
print("model has {} trainable parameters".format(num_params))
# model = nn.DataParallel(model)
model.to(device)
model.cuda()
class_weights = torch.Tensor((1, 1)).float()
class_weights = class_weights.to(device)
loss_segmentation = nn.CrossEntropyLoss(class_weights)
# loss_convexity = loss.ConvexShapeLoss(height=456,width=608,device=device)
optimizer = optim.Adam(model.parameters(),
lr=config['testing']['learning_rate'],
batch_size=1,
num_workers=args.nThreads,
collate_fn=csv_collator)
# Array with [height, width] of the new size
resized_size = np.array([args.height, args.width])
# Loss function
criterion_training = losses.WeightedHausdorffDistance(
resized_height=args.height,
resized_width=args.width,
return_2_terms=True,
device=device)
# Restore saved checkpoint (model weights)
with peter("Loading checkpoint"):
# Pretrained models that come with this package
if args.model == 'unet_256x256_sorghum':
args.model = os.path.join(os.path.dirname(os.path.realpath(__file__)),
'checkpoints', 'unet_256x256_sorghum.ckpt')
if os.path.isfile(args.model):
if args.cuda:
checkpoint = torch.load(args.model)
else:
checkpoint = torch.load(args.model,
map_location=lambda storage, loc: storage)
# Model
if args.n_points is None:
if 'n_points' not in checkpoint:
# Model will also estimate # of points
