def _load(self, filePath):
checkpoint = torch.load(filePath)
model = MyModel(
device, checkpoint['inputSize'], checkpoint['gatedCnnOutputSize'],
checkpoint['gatedCnnStride1'], checkpoint['gatedCnnStride2'],
checkpoint['gatedCnnKernel1'], checkpoint['gatedCnnKernel2'],
checkpoint['lstmLayer'], checkpoint['lstmHiddenSize'],
checkpoint['fcOutputSize'], checkpoint['dropout'])
model.load_state_dict(checkpoint['stateDict'])
model.eval()
if self.device.type == 'cpu':
model.cpu()
else:
model.cuda(device=self.device)
return model
def main(args):
# Create model directory
full_model_path = args.model_path
# Image preprocessing
transform = transforms.Compose(
[transforms.Scale(args.crop_size),
transforms.ToTensor()])
#transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
# Load vocabulary wrapper.
vocab = build_vocab(args.image_dir, 1, None)
# Build data loader
data_loader = get_loader(args.image_dir,
vocab,
transform,
args.batch_size,
shuffle=True,
num_workers=2)
data_set = ProcessingDataset(root=args.image_dir,
vocab=vocab,
transform=transform)
train_loader = torch.utils.data.DataLoader(data_set, collate_fn=collate_fn)
train_size = len(train_loader)
# Build the models
model = MyModel(args.embed_size, args.hidden_size, len(vocab), vocab)
print(model)
if torch.cuda.is_available():
model.cuda()
# Loss and Optimizer
criterion = nn.MSELoss()
#criterion = nn.SmoothL1Loss()
params = list(model.parameters())
optimizer = torch.optim.Adam(params, lr=args.learning_rate)
# Train the Models
total_step = len(data_loader)
for epoch in range(args.num_epochs):
for i, (images, captions, lengths) in enumerate(data_loader):
model.train()
image_ts = to_var(images)
captions = to_var(captions)
count = images.size()[0]
# Forward, Backward and Optimize
model.zero_grad()
outputs = model(captions, lengths)
loss = criterion(outputs, image_ts)
loss.backward()
optimizer.step()
correct = outputs.data.eq(image_ts.data).sum()
accuracy = 100. * correct / count
# Print log info
if i % args.log_step == 0:
#print("i "+str(i))
#torch.set_printoptions(profile="full")
for ii, t in enumerate(outputs):
result = transforms.ToPILImage()(t.data.cpu())
result.save("./results/" + str(i) + "_" + str(ii) + ".png")
origin = transforms.ToPILImage()(image_ts[ii].data.cpu())
origin.save("./results/" + str(i) + "_" + str(ii) +
"target.png")
with open(
("./results/" + str(i) + "_" + str(ii) + "_diff.txt"),
'w') as f:
f.write(str(torch.abs(t - image_ts[ii]).sum()))
print(
'Epoch [%d/%d], Step [%d/%d], Loss: %.4f, accuracy: %2.2f Perplexity: %5.4f'
% (epoch, args.num_epochs, i, total_step, loss.data[0],
accuracy, np.exp(loss.data[0])))
# Save the models
if (i + 1) % args.save_step == 0:
torch.save(
model.state_dict(),
os.path.join(full_model_path,
'decoder-%d-%d.pkl' % (epoch + 1, i + 1)))
torch.save(
model.state_dict(),
os.path.join(full_model_path,
'decoder-%d-%d.pkl' % (epoch + 1, i + 1)))
class RunGC(object):
@classmethod
def get_option_spec(cls):
spec = PyOptionSpec()
elf.saveDefaultOptionsToArgs("", spec)
elf.saveDefaultNetOptionsToArgs("", spec)
spec.addIntOption(
'gpu',
'GPU id to use',
-1)
spec.addStrListOption(
"parsed_args",
"dummy option",
[])
return spec
@auto_import_options
def __init__(self, option_map):
self.option_map = option_map
def initialize(self):
opt = elf.Options()
net_opt = elf.NetOptions()
opt.loadFromArgs("", self.option_map.getOptionSpec())
net_opt.loadFromArgs("", self.option_map.getOptionSpec())
self.rs = elf.RemoteServers(elf.getNetOptions(opt, net_opt), ["actor", "train"])
GC = elf.BatchReceiver(opt, self.rs)
GC.setMode(elf.RECV_ENTRY)
batchsize = opt.batchsize
print("Batchsize: %d" % batchsize)
width = 210 // 2
height = 160 // 2
T = 6
num_action = 4
spec = {}
spec["actor"] = dict(
input=dict(s=("float", (3, height, width))),
reply=dict(a=("int32_t", 1), pi=("float", num_action), V=("float", 1))
)
'''
spec["train"] = dict(
input=dict(s_=(T, 3, height, width), r_=(T, 1), a_=(T, 1), pi_=(T, num_action), V_=(T, 1)),
)
'''
e = GC.getExtractor()
desc = allocExtractor(e, batchsize, spec)
params = {
"input_dim" : width * height * 3,
"num_action" : 4
}
print("Init GC Wrapper")
has_gpu = self.options.gpu is not None and self.options.gpu >= 0
self.wrapper = GCWrapper(
GC, None, batchsize, desc, num_recv=1, default_gpu=(self.options.gpu if has_gpu else None),
use_numpy=False, params=params)
# wrapper.reg_callback("train", self.on_train)
self.wrapper.reg_callback("actor", self.on_actor)
self.model = MyModel(params)
if has_gpu:
self.model.cuda(self.options.gpu)
# self.optim = torch.optimi.Adam(self.model.parameters())
self.n = 0
def on_actor(self, batch):
res = self.model(batch)
m = torch.distributions.Categorical(res["pi"].data)
self.n += 1
if self.n == 20:
# gives a single float value
#print(psutil.cpu_percent())
# gives an object with many fields
#print(psutil.virtual_memory())
self.n = 0
return dict(a=m.sample(), pi=res["pi"].data, V=res["V"].data)
def on_train(self, batch):
pass
download=True)
test_data_loader = data.DataLoader(test_dataset,
batch_size=opt.batch_size,
shuffle=False,
num_workers=opt.num_workers)
# instantiate network (which has been imported from *networks.py*)
net = MyModel()
# create losses (criterion in pytorch)
criterion_CE = nn.CrossEntropyLoss()
# if running on GPU and we want to use cuda move model there
use_cuda = torch.cuda.is_available()
if use_cuda:
net = net.cuda()
# create optimizers
optim = torch.optim.Adam(net.parameters(), lr=opt.lr)
# load checkpoint if needed/ wanted
start_n_iter = 0
start_epoch = 0
if opt.resume:
ckpt = load_checkpoint(opt.path_to_checkpoint
) # custom method for loading last checkpoint
net.load_state_dict(ckpt['net'])
start_epoch = ckpt['epoch']
start_n_iter = ckpt['n_iter']
optim.load_state_dict(ckpt['optim'])
print("last checkpoint restored")
return mae, r
if __name__ == "__main__":
opt = get_args()
if opt.gpu:
torch.cuda.manual_seed(0)
else:
torch.manual_seed(0)
trainloader, devloader, testloader = get_data(opt)
model = MyModel(opt)
if opt.gpu:
model = nn.DataParallel(model)
model.cuda()
optimizer = AdamW(model.parameters(), lr=opt.lr)
log = open(opt.log_path, "a")
max_r = 0
for epoch in range(opt.nepoch):
train_loss = train(model, trainloader, optimizer, opt)
dev_mae, dev_r = test(model, devloader, opt)
print("epoch: %d train_loss: %f dev_mae: %.3f dev_r: %.3f" %
(epoch, train_loss, dev_mae, dev_r))
print("epoch: %d train_loss: %f dev_mae: %.3f dev_r: %.3f" %
(epoch, train_loss, dev_mae, dev_r),
file=log)
if dev_r > max_r:
