def main(args):
net = resnet18()
optimizer = torch.optim.Adam(net.parameters(), lr=args.lr)
# load data
train_iter, test_iter = mlutils.load_data_fashion_mnist(batch_size=args.batch_size, resize=96)
# train
mlutils.train(net, train_iter, test_iter, args.batch_size, optimizer, args.num_epochs)
def main(args):
# The original VGG network has 5 convolutional blocks.
# The first two blocks have one convolutional layer.
# The latter three blocks contain two convolutional layers.
conv_arch = ((1, 64), (1, 128), (2, 256), (2, 512), (2, 512))
# The parameters of VGG-11 are big, use a ratio to reduce the network size by dividing a ratio on the output channel number.
ratio = args.ratio
small_conv_arch = [(pair[0], pair[1] // ratio) for pair in conv_arch]
net = vgg(small_conv_arch)
optimizer = torch.optim.Adam(net.parameters(), lr=args.lr)
# load data
train_iter, test_iter = mlutils.load_data_fashion_mnist(batch_size=args.batch_size, resize=224)
# train
train(net, train_iter, test_iter, args.batch_size, optimizer, args.num_epochs)