def main():
# Device configuration, check cuda availability
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# Hyper parameters
BATCH_SIZE = 128
EPOCHS_NUM = 50
LR = 0.001
GAMMA = 0.96
transform = transforms.Compose([
# shift by 2 pixels in either direction with zero padding.
transforms.RandomCrop(28, padding=2),
transforms.ToTensor(),
transforms.Normalize((0.1307, ), (0.3081, ))
])
train_dataset = torchvision.datasets.MNIST(root=op.join(
sys.path[0], 'data/'),
train=True,
transform=transform,
download=True)
test_dataset = torchvision.datasets.MNIST(root=op.join(
sys.path[0], 'data/'),
train=False,
transform=transform)
# Data loader
train_loader = Data.DataLoader(dataset=train_dataset,
batch_size=BATCH_SIZE,
num_workers=4,
shuffle=True)
test_loader = Data.DataLoader(dataset=test_dataset,
batch_size=BATCH_SIZE,
num_workers=4,
shuffle=True)
# Load model
model = CapsNet().to(device)
# Loss and optimizer
criterion = CapsuleLoss()
optimizer = optim.Adam(model.parameters(), lr=LR)
scheduler = optim.lr_scheduler.ExponentialLR(optimizer, gamma=GAMMA)
# Train and test model
train(model, device, train_loader, criterion, optimizer, scheduler,
EPOCHS_NUM)
test(model, device, test_loader)
# Save model
torch.save(model, op.join(sys.path[0], 'model/mnist_capsnet.pt'))
loss = capsule_net.loss(data, output, target, reconstructions)
test_loss += loss.data[0]
correct += sum(
np.argmax(masked.data.cpu().numpy(), 1) == np.argmax(
target.data.cpu().numpy(), 1))
tqdm.write("Epoch: [{}/{}], test accuracy: {:.6f}, loss: {:.6f}".format(
epoch, N_EPOCHS, correct / len(test_loader.dataset),
test_loss / len(test_loader)))
if __name__ == '__main__':
torch.manual_seed(1)
dataset = 'cifar10'
# dataset = 'mnist'
config = Config(dataset)
mnist = Dataset(dataset, BATCH_SIZE)
capsule_net = CapsNet(config)
capsule_net = torch.nn.DataParallel(capsule_net)
if USE_CUDA:
capsule_net = capsule_net.cuda()
capsule_net = capsule_net.module
optimizer = torch.optim.Adam(capsule_net.parameters())
for e in range(1, N_EPOCHS + 1):
train(capsule_net, optimizer, mnist.train_loader, e)
test(capsule_net, mnist.test_loader, e)
args.batch_size,
args.use_gpu)
model = CapsNet(n_conv_channel=256,
n_primary_caps=8,
primary_cap_size=1152,
output_unit_size=16,
n_routing_iter=3)
# load state from past runs
if args.load_checkpoint != '':
model.load_state_dict(torch.load(args.load_checkpoint))
# move to GPU
model = model.cuda() if args.use_gpu else model
optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
# setup decoder for training
decoder = Decoder()
decoder = decoder.cuda() if args.use_gpu else decoder
decoder_optim = torch.optim.Adam(decoder.parameters(), lr=0.001)
# use decaying learning rate
scheduler = lr_scheduler.ExponentialLR(decoder_optim, 0.5)
for epoch in range(1, args.epochs + 1):
train(epoch, model, trainloader, optimizer, decoder, decoder_optim)
test(epoch, model, testloader)
scheduler.step()
if args.checkpoint_interval > 0:
from torch.utils.data.dataloader import DataLoader
from torchvision import datasets
from capsnet import CapsNet
from functions import DigitMarginLoss
from functions import accuracy
train_loader = DataLoader(datasets.MNIST('data', train=True, download=True, transform=transforms.Compose([
# transforms.RandomShift(2),
transforms.ToTensor()])), batch_size=1, shuffle=True)
test_loader = DataLoader(datasets.MNIST('data', train=False, transform=transforms.Compose([
transforms.ToTensor()])), batch_size=1)
model = CapsNet()
optimizer = optim.Adam(model.parameters())
margin_loss = DigitMarginLoss()
reconstruction_loss = torch.nn.MSELoss(size_average=False)
model.train()
for epoch in range(1, 11):
epoch_tot_loss = 0
epoch_tot_acc = 0
for batch, (data, target) in enumerate(train_loader, 1):
data = Variable(data)
target = Variable(target)
digit_caps, reconstruction = model(data, target)
loss = margin_loss(digit_caps, target) + 0.0005 * reconstruction_loss(reconstruction, data.view(-1))
epoch_tot_loss += loss
"""
#imgs = imgs[:2]
stat = []
model = CapsNet(args,
len(train_dataset) // (2 * args.batch_size) + 3, stat)
use_cuda = not args.disable_cuda and torch.cuda.is_available()
if use_cuda:
model.cuda()
imgs = imgs.cuda()
if args.jit:
model = torch.jit.trace(model, (imgs), check_inputs=[(imgs)])
else:
model(imgs)
print("# model parameters:",
sum(param.numel() for param in model.parameters()))
"""
Construct optimizer, scheduler, and loss functions
"""
optimizer = torch.optim.Adam(model.parameters(),
lr=args.lr if args.lr else 1e-3,
betas=(0.9, 0.999))
#optimizer = torch.optim.SGD(model.parameters(), lr=args.lr if args.lr else 1e-3, momentum=0.9, weight_decay=0.0005)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
optimizer,
mode='min',
factor=0.5,
patience=args.patience,
threshold=1e-6,
verbose=True)
def main():
# Load model
model = CapsNet().to(device)
criterion = CapsuleLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-3)
scheduler = optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.96)
# Load data
transform = transforms.Compose([
# shift by 2 pixels in either direction with zero padding.
transforms.RandomCrop(28, padding=2),
transforms.ToTensor(),
transforms.Normalize((0.1307, ), (0.3081, )),
])
DATA_PATH = "./data"
BATCH_SIZE = 128
train_loader = DataLoader(
dataset=MNIST(root=DATA_PATH,
download=True,
train=True,
transform=transform),
batch_size=BATCH_SIZE,
num_workers=4,
shuffle=True,
)
test_loader = DataLoader(
dataset=MNIST(root=DATA_PATH,
download=True,
train=False,
transform=transform),
batch_size=BATCH_SIZE,
num_workers=4,
shuffle=True,
)
# Train
EPOCHES = 50
model.train()
for ep in range(EPOCHES):
batch_id = 1
correct, total, total_loss = 0, 0, 0.0
for images, labels in train_loader:
optimizer.zero_grad()
images = images.to(device)
labels = torch.eye(10).index_select(dim=0, index=labels).to(device)
logits, reconstruction = model(images)
# Compute loss & accuracy
loss = criterion(images, labels, logits, reconstruction)
correct += torch.sum(
torch.argmax(logits, dim=1) == torch.argmax(labels,
dim=1)).item()
total += len(labels)
accuracy = correct / total
total_loss += loss
loss.backward()
optimizer.step()
print("Epoch {}, batch {}, loss: {}, accuracy: {}".format(
ep + 1, batch_id, total_loss / batch_id, accuracy))
batch_id += 1
scheduler.step(ep)
print("Total loss for epoch {}: {}".format(ep + 1, total_loss))
# Eval
model.eval()
correct, total = 0, 0
for images, labels in test_loader:
# Add channels = 1
images = images.to(device)
# Categogrical encoding
labels = torch.eye(10).index_select(dim=0, index=labels).to(device)
logits, reconstructions = model(images)
pred_labels = torch.argmax(logits, dim=1)
correct += torch.sum(pred_labels == torch.argmax(labels, dim=1)).item()
total += len(labels)
print("Accuracy: {}".format(correct / total))
# Save model
torch.save(
model.state_dict(),
"./model/capsnet_ep{}_acc{}.pt".format(EPOCHES, correct / total),
)
DEVICE = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
DATASET_CONFIG = MNIST
dataset_transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
])
train_dataset = datasets.MNIST('./data',
train=True,
download=True,
transform=dataset_transform)
test_dataset = datasets.MNIST('./data',
train=False,
download=True,
transform=dataset_transform)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False)
torch.manual_seed(1)
capsule_net = CapsNet(**DATASET_CONFIG)
capsule_net.to(DEVICE)
optimizer = torch.optim.Adam(capsule_net.parameters(), lr=LEARNING_RATE)
for e in range(1, 1 + EPOCHS):
train(capsule_net, optimizer, train_loader, e, device=DEVICE)
evaluate(capsule_net, test_loader, e, device=DEVICE)
from utils import accuracy
train_loader = torch.utils.data.DataLoader(
datasets.MNIST('data', train=True, download=True, transform=transforms.Compose([
# transforms.RandomShift(2),
transforms.ToTensor()
])), shuffle=True)
test_loader = torch.utils.data.DataLoader(
datasets.MNIST('data', train=False, transform=transforms.Compose([
transforms.ToTensor(),
])))
model = CapsNet()
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
margin_loss = DigitMarginLoss()
reconstruction_loss = torch.nn.MSELoss(size_average=False)
model.train()
for epoch in range(1, 11):
epoch_tot_loss = 0
epoch_tot_acc = 0
for batch, (input, target) in enumerate(train_loader, 1):
input = Variable(input)
target = Variable(target)
digit_caps, reconstruction = model(input, target)
loss = margin_loss(digit_caps, target) + 0.0005 * reconstruction_loss(reconstruction, input.view(-1))
epoch_tot_loss += loss
from torch.autograd import Variable
import torch.utils.data as Data
# train_data:[17023,6000,1],train_tag:[17023] numpy类型
train_data,train_tag=get_train_data()
# test_data,test_tag=get_test_data()
net=CapsNet()
# 数据集加载
# 数据转换成tensor 并[17023,6000,1]=>[17023,1,6000]
train_data=torch.FloatTensor(train_data).permute(0,2,1)
train_tag=torch.LongTensor(train_tag)
train_set=Data.TensorDataset(train_data,train_tag)
train_loader=Data.DataLoader(dataset=train_set,batch_size=32,shuffle=True)
# 优化器设置
optimizer=torch.optim.Adam(net.parameters(),lr=0.001)
loss_func=torch.nn.CrossEntropyLoss()
# train...
print('开始训练：')
accc=[]
epoch=50
maxnum=0
for i in range(epoch):
for j,(x,y) in enumerate(train_loader):
x,y=Variable(x),Variable(y)
out=net(x)
loss=loss_func(out,y)
optimizer.zero_grad()
loss.backward()
