joint_transforms.RandomRotate(),
joint_transforms.Resize((args['scale'], args['scale']))
])
val_joint_transform = joint_transforms.Compose(
[joint_transforms.Resize((args['scale'], args['scale']))])
img_transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406],
[0.229, 0.224, 0.225]) # maybe can optimized.
])
target_transform = transforms.ToTensor()
# Prepare Data Set.
train_set = ImageFolder(msd_training_root, joint_transform, img_transform,
target_transform)
print("Train set: {}".format(train_set.__len__()))
train_loader = DataLoader(train_set,
batch_size=args['train_batch_size'],
num_workers=0,
shuffle=True)
val_set = ImageFolder(msd_testing_root, val_joint_transform, img_transform,
target_transform)
print("Validation Set: {}".format(val_set.__len__()))
val_loader = DataLoader(val_set,
batch_size=args['val_batch_size'],
num_workers=8,
shuffle=False)
bce = nn.BCEWithLogitsLoss().cuda(device_ids[0])
# Transform Data.
joint_transform = joint_transforms.Compose([
joint_transforms.RandomRotate(),
joint_transforms.Resize((args['scale'], args['scale']))
])
img_transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406],
[0.229, 0.224, 0.225]) # maybe can optimized.
])
target_transform = transforms.ToTensor()
# Prepare Data Set.
train_set = ImageFolder(msd_training_root, joint_transform, img_transform,
target_transform)
print("Train set: {}".format(train_set.__len__()))
train_loader = DataLoader(train_set,
batch_size=args['train_batch_size'],
num_workers=0,
shuffle=True)
def main():
print(args)
print(exp_name)
net = BASE3(backbone_path).cuda(device_ids[0]).train()
if args['add_graph']:
writer.add_graph(net,
input_to_model=torch.rand(
args['train_batch_size'], 3, args['scale'],