from model import Model
from options import TrainOptions
import torch
from torchvision.transforms import *
import torch.optim as optim
import numpy as np
from visualizer import Visualizer
import os
# Get the Hyperparaeters
opt = TrainOptions().parse()
import matplotlib.pyplot as plt
sample_dataset = DataSet(opt,"./dataset/T1-train/img/", "./dataset/T1-train/GT/T1-GT.xml","./dataset/T2-Train/img/" )
train_sampler,val_sampler = create_samplers(sample_dataset.__len__(),opt.split_ratio)
sample_loader = torch.utils.data.DataLoader(sample_dataset,sampler=train_sampler,batch_size=opt.batch_size,num_workers=opt.workers)
sample_val_loader = torch.utils.data.DataLoader(sample_dataset,sampler=val_sampler,batch_size=opt.val_batch_size,num_workers=opt.workers//5,shuffle=False)
# Check if gpu available or not
device = torch.device("cuda" if (torch.cuda.is_available() and opt.use_gpu) else "cpu")
opt.device = device
# Load the model and send it to gpu
model = Model(opt)
model = model.to(device)
if opt.use_gpu:
model = torch.nn.DataParallel(model, device_ids=opt.gpus)
# Print our model
print('------------ Model -------------')
# Main module to train the model, load the data,
# do gradient descent etc. followed by saving our model for later testing
from dataloader import DataSet, create_samplers
from model import Model
from visualizer import Visualizer
from options import TrainOptions
import torch
from torchvision.transforms import *
import torch.optim as optim
import numpy as np
import os
# Get the Hyperparaeters
opt = TrainOptions().parse()
sample_dataset = DataSet(opt, "/media/shubh/PranayHDD/Kinect/")
train_sampler, val_sampler = create_samplers(sample_dataset.__len__(),
opt.split_ratio)
data_loader = torch.utils.data.DataLoader(sample_dataset,
sampler=train_sampler,
batch_size=opt.batch_size,
num_workers=opt.num_workers)
data_val_loader = torch.utils.data.DataLoader(sample_dataset,
sampler=val_sampler,
batch_size=opt.val_batch_size,
num_workers=0,
shuffle=False)
# Check if gpu available or not
device = torch.device("cuda" if (
torch.cuda.is_available() and opt.use_gpu) else "cpu")
opt.device = device