dataset = Rand_num()
sampler = RandomSampler(dataset)
loader = DataLoader(dataset,
batch_size=20,
sampler=sampler,
shuffle=False,
num_workers=1,
drop_last=True)
net = AlexNet(3)
#net.load_state_dict(torch.load(SAVE_PATH))
net.cuda()
optimizer = optim.Adam(net.parameters(), lr=0.001)
for epoch in range(10000):
for i, data in enumerate(loader, 0):
net.zero_grad()
video, labels = data
video = video.view(-1, 3, 227, 227)
labels = labels.view(-1, 3)
labels = torch.squeeze(Variable(labels.float().cuda()))
video = torch.squeeze(Variable((video.float() / 256).cuda()))
net.train()
outputs = net.forward(video)
loss = lossfunction(outputs, labels)
loss.backward()
optimizer.step()
if i == 0:
torch.save(net.state_dict(), SAVE_PATH)
print(loss)
logger.scalar_summary('loss', loss.data.cpu().numpy(), epoch)
