# Train
net.train()
cont = 0
loss_train = 0.0
grads_loss = 0.0
RGB_drops = np.array([1] *
n_epoch) # + list(range(5)) + [5]*(n_epoch-10))/5
#RGB_drops = np.array([0]*n_epoch + list(range(5)) + [5]*(n_epoch-10))/5
# flip
#RGB_drops = RGB_drops[::-1]
# Transforms train
train_trans = make_train_transforms(drop_p=RGB_drops[epoch])
test_trans = make_test_transforms(drop_p=RGB_drops[epoch])
# Create datasets
dataset = NYUDataset(train_depths, transforms=train_trans)
dataset_val = NYUDataset(depths_list['val'], transforms=test_trans)
training_generator = data.DataLoader(dataset, **params)
val_generator = data.DataLoader(dataset_val, **params_test)
for _i, (depths, rgbs, filename) in enumerate(training_generator):
#cont+=1
iter_train += 1
# Get items from generator
inputs, outputs = rgbs.cuda(), depths.cuda()
#print(torch.max(outputs.view(input.size(0), -1)))
# Clean grads
optimizer_ft.zero_grad()
CLAHE(clip_limit=2),
IAASharpen(),
IAAEmboss(),
RandomBrightnessContrast(),
], p=0.3),
HueSaturationValue(p=0.3),
], p=p)
augm = strong_aug(0.9)
depths = ['../sample_images/classroom_000310depth.png','../sample_images/classroom_000350depth.png',
'../sample_images/classroom_000329depth.png']
# Instantiate a model and dataset
dataset = NYUDataset(depths, is_train = False, transforms= augm)
# Parameters
params = {'batch_size': 16 ,
'shuffle': True,
'num_workers': 12,
'pin_memory': True}
params_test = {'batch_size': 16 ,
'shuffle': False,
'num_workers': 12,
'pin_memory': True}
training_generator = data.DataLoader(dataset,**params)
val_generator = data.DataLoader(dataset_val,**params_test)
model = RGBDepth_Depth()