writer.add_image('Train.3.Diff', colorize(vutils.make_grid(torch.abs(output-depth).data, nrow=6, normalize=False)), epoch)
del image
del depth
del output
traincsv=pd.read_csv('./content/data/diml_outdoor_train.csv')
traincsv = traincsv.values.tolist()
traincsv = shuffle(traincsv, random_state=2)
#display a sample set of image and depth image
depth_dataset = DepthDataset(traincsv=traincsv,root_dir='./content/')
fig = plt.figure()
len(depth_dataset)
model = Model().cpu()
if torch.cuda.device_count() > 1:
print("Let's use", torch.cuda.device_count(), "GPUs!")
# dim = 0 [30, xxx] -> [10, ...], [10, ...], [10, ...] on 3 GPUs
model = nn.DataParallel(model)
#load trained model if needed
#model.load_state_dict(torch.load('/workspace/1.pth'))
print('Model created.')
epochs=50
lr=0.0001
batch_size=64
depth_dataset = DepthDataset(traincsv=traincsv, root_dir='./content/',
transform=transforms.Compose([Augmentation(0.5),ToTensor()]))
train_loader=DataLoader(depth_dataset, batch_size, shuffle=True)
start_time = time.time()
depth_dataset = DepthDataset(root_dir=data,
transform=transforms.Compose([ToTensor()]))
train_loader = torch.utils.data.DataLoader(depth_dataset, batchSize)
dataiter = iter(train_loader)
images = dataiter.next()
print("\n Time taken to load Images: %s " % (time.time() - start_time))
print("\n Test Dataset Shape: {shape}".format(shape=np.shape(depth_dataset)))
# ### Importing the Model
from Mobile_model import Model
model = Model().cuda()
model = nn.DataParallel(model)
# Import the Pre-trained Model
model.load_state_dict(torch.load(pretrained_path))
print("\n Loaded MobileNet U-Net Weights successfully\n")
model.eval()
# ### Model Variables (state_dict)
# print("\n\nModel's state_dict:\n\n")
# for param_tensor in model.state_dict():
# print(param_tensor, "\t", model.state_dict()[param_tensor].size())
help='Image size of network input')
parser.add_argument('--data_dir',
default='comarision_datasets\input',
type=str,
help='Data path')
parser.add_argument(
'--result_dir',
default='demo_results',
type=str,
help='Directory for saving results, default: demo_results')
parser.add_argument('--gpu_id',
default=0,
type=int,
help='GPU id, default:0')
args = parser.parse_args()
if not os.path.exists(args.result_dir):
os.makedirs(args.result_dir)
gpu_id = args.gpu_id
torch.cuda.device(gpu_id)
net = Model().cuda()
net.load_state_dict(torch.load('weights_model.pth'))
net.eval()
print('Begin to test ...')
with torch.no_grad():
demo(net, args)
print('Finished!')