def setUp(self, mock):
# prepare mock.
mock.return_value = ['image1.png,108.0,50.0,1.0,148.0,180.0,0.0,148.0,180.0,1.0\n',
'image2.png,40.0,50.0,1.0,160.0,180.0,1.0,160.0,180.0,0.0\n']
# set up.
self.path = 'test_data'
self.dataset = PoseDataset(
self.path,
input_transform=transforms.Compose([
transforms.ToTensor(),
RandomNoise()]),
output_transform=Scale(),
transform=Crop(data_augmentation=True))
def __init__(self, Nj, gpu, model_file, filename):
# validate arguments.
self.gpu = (gpu >= 0)
if self.gpu and not torch.cuda.is_available():
raise GPUNotFoundError('GPU is not found.')
# initialize model to estimate.
self.model = AlexNet(Nj)
self.model.load_state_dict(torch.load(model_file))
# prepare gpu.
if self.gpu:
self.model.cuda()
# load dataset to estimate.
self.dataset = PoseDataset(filename,
input_transform=transforms.Compose(
[transforms.ToTensor(),
RandomNoise()]),
output_transform=Scale(),
transform=Crop())
def start(self):
""" Train pose net. """
# set random seed.
if self.seed is not None:
random.seed(self.seed)
torch.manual_seed(self.seed)
if self.gpu:
torch.cuda.manual_seed(self.seed)
# initialize model to train.
model = AlexNet(self.Nj)
if self.resume_model:
model.load_state_dict(torch.load(self.resume_model))
# prepare gpu.
if self.gpu:
model.cuda()
# load the datasets.
input_transforms = [transforms.ToTensor()]
if self.data_augmentation:
input_transforms.append(RandomNoise())
train = PoseDataset(
self.train,
input_transform=transforms.Compose(input_transforms),
output_transform=Scale(),
transform=Crop(data_augmentation=self.data_augmentation))
val = PoseDataset(self.val,
input_transform=transforms.Compose(
[transforms.ToTensor()]),
output_transform=Scale(),
transform=Crop(data_augmentation=False))
# training/validation iterators.
train_iter = torch.utils.data.DataLoader(train,
batch_size=self.batchsize,
shuffle=True)
val_iter = torch.utils.data.DataLoader(val,
batch_size=self.batchsize,
shuffle=False)
# set up an optimizer.
optimizer = self._get_optimizer(model)
if self.resume_opt:
optimizer.load_state_dict(torch.load(self.resume_opt))
# set intervals.
val_interval = 10
resume_interval = self.epoch / 10
log_interval = 10
# set logger and start epoch.
logger = TrainLogger(os.path.join(self.out, 'pytorch'))
start_epoch = 0
if self.resume:
resume = torch.load(self.resume)
start_epoch = resume['epoch']
logger.load_state_dict(resume['logger'])
# start training.
start_time = time.time()
for epoch in trange(start_epoch,
self.epoch,
initial=start_epoch,
total=self.epoch,
desc=' total'):
self._train(model, optimizer, train_iter, log_interval, logger,
start_time)
if (epoch + 1) % val_interval == 0:
self._test(model, val_iter, logger, start_time)
if (epoch + 1) % resume_interval == 0:
self._checkpoint(epoch, model, optimizer, logger)
def setUp(self):
self.transform = RandomNoise()
def __init__(self,
Nj,
NN,
gpu,
model_file,
filename,
Dataset3D,
isEval=True):
# validate arguments.
self.gpu = (gpu >= 0)
self.NN = NN
if self.gpu and not torch.cuda.is_available():
raise GPUNotFoundError('GPU is not found.')
# initialize model to estimate.
if self.NN == "MobileNet_":
self.model = MobileNet_()
elif self.NN == "MobileNet__":
self.model = MobileNet__()
elif self.NN == "MobileNet_3":
self.model = MobileNet_3()
elif self.NN == "MobileNet_4":
self.model = MobileNet_4()
elif self.NN == "MobileNet___":
self.model = MobileNet___()
elif self.NN == "MobileNet":
self.model = MobileNet()
elif self.NN == "MnasNet":
self.model = MnasNet()
elif self.NN == "MnasNet_":
self.model = MnasNet_()
elif self.NN == "MnasNet56_":
self.model = MnasNet56_()
elif self.NN == "AlexNet":
self.model = AlexNet(Nj)
elif self.NN == "MobileNet3D":
self.model = MobileNet3D()
elif self.NN == "MnasNet3D":
self.model = MnasNet3D()
elif self.NN == "MobileNet3D2":
self.model = MobileNet3D2()
else:
self.model = Resnet()
self.model.load_state_dict(torch.load(model_file))
if isEval == True:
self.model.eval()
# prepare gpu.
if self.gpu:
self.model.cuda()
# load dataset to estimate.
if Dataset3D:
'''
self.dataset = PoseDataset3D(
filename,
input_transform=transforms.ToTensor())
'''
self.dataset = PoseDataset3D(filename,
input_transform=transforms.Compose([
transforms.ToTensor(),
RandomNoise()
]))
else:
self.dataset = PoseDataset(filename,
input_transform=transforms.Compose([
transforms.ToTensor(),
RandomNoise()
]),
output_transform=Scale(),
transform=Crop(data_augmentation=False))