def run(args):
if args.use_gpu:
try:
output = subprocess.check_output("free-gpu", shell=True)
output = int(output.decode('utf-8'))
gpu = output
print("claiming gpu {}".format(gpu))
torch.cuda.set_device(int(gpu))
a = torch.tensor([1]).cuda()
device = int(gpu)
except (IndexError, subprocess.CalledProcessError) as e:
device = None
else:
device = None
print(f"device is : {device}")
train_loader = DataLoader("../../data/raw/xkcd_colordata",
"../../data/raw",
"train",
device=device)
dev_generator = DataLoader("../../data/raw/xkcd_colordata",
"../../data/raw/",
"dev",
device=device)
model = LinearRegression(128, 3)
if device is not None:
model = model.cuda(device)
criterion = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=args.lr)
for i in range(args.n_epochs):
for batch in train_loader:
X, Y = prepare_batch(batch, device)
optimizer.zero_grad()
Y_pred = model.forward(X)
loss = criterion(Y_pred, Y)
loss.backward()
optimizer.step()
print(f"epoch: {i}, loss: {loss.data[0]}")
torch.save(model, "../../models/baseline/{args.lr}_linear_{i}")
