def train(model, data, device, maxepoch):
torch.manual_seed(100)
train_loader, test_loader = feed_dataset(data, 'deeprobust/image/data')
if (model == 'CNN'):
import deeprobust.image.netmodels.CNN as MODEL
from deeprobust.image.netmodels.CNN import Net
train_net = Net().to(device)
elif (model == 'ResNet18'):
import deeprobust.image.netmodels.resnet as MODEL
train_net = MODEL.ResNet18().to(device)
elif (model == 'ResNet34'):
import deeprobust.image.netmodels.resnet as MODEL
train_net = MODEL.ResNet34().to(device)
elif (model == 'ResNet50'):
import deeprobust.image.netmodels.resnet as MODEL
train_net = MODEL.ResNet50().to(device)
optimizer = optim.SGD(train_net.parameters(), lr=0.01, momentum=0.5)
save_model = True
for epoch in range(1, maxepoch + 1): ## 5 batches
print(epoch)
MODEL.train(train_net, device, train_loader, optimizer, epoch)
MODEL.test(train_net, device, test_loader)
if (save_model and epoch % 10 == 0):
if os.path.isdir('./trained_models/'):
print('Save model.')
torch.save(
train_net.state_dict(), './trained_models/' + data + "_" +
model + "_epoch_" + str(epoch) + ".pt")
else:
os.mkdir('./trained_models/')
print('Make directory and save model.')
torch.save(
train_net.state_dict(), './trained_models/' + data + "_" +
model + "_epoch_" + str(epoch) + ".pt")
def generate(self, train_loader, test_loader, **kwargs):
self.parse_params(**kwargs)
# settings
if not os.path.exists(model_dir):
os.makedirs(model_dir)
use_cuda = not self.no_cuda and torch.cuda.is_available()
torch.manual_seed(self.seed)
loader_kwargs = {'num_workers': 1, 'pin_memory': True} if use_cuda else {}
# setup data loader
train_loader = torch.utils.data.DataLoader(
datasets.MNIST('../data', train=True, download=True,
transform=transforms.ToTensor()),
batch_size=args.batch_size, shuffle=True, **loader_kwargs)
test_loader = torch.utils.data.DataLoader(
datasets.MNIST('../data', train=False,
transform=transforms.ToTensor()),
batch_size=args.test_batch_size, shuffle=True, **loader_kwargs)
# init model, Net() can be also used here for training
model = Net().to(device)
optimizer = optim.SGD(model.parameters(), lr = self.lr, momentum=self.momentum)
for epoch in range(1, self.epochs + 1):
# adjust learning rate for SGD
optimizer = adjust_learning_rate(optimizer, epoch, self.lr)
# adversarial training
train(args, model, device, train_loader, optimizer, epoch)
# evaluation on natural examples
print('================================================================')
eval_train(model, device, train_loader)
eval_test(model, device, test_loader)
print('================================================================')
# save checkpoint
if epoch % self.save_freq == 0:
torch.save(model.state_dict(),
os.path.join(model_dir, 'trade_model-nn-epoch{}.pt'.format(epoch)))
torch.save(optimizer.state_dict(),
os.path.join(model_dir, 'opt-nn-checkpoint_epoch{}.tar'.format(epoch)))
shuffle=True)
test_loader = torch.utils.data.DataLoader(datasets.MNIST(
'deeprobust/image/data',
train=False,
transform=transforms.Compose([transforms.ToTensor()])),
batch_size=1000,
shuffle=True)
#ipdb.set_trace()
#TODO: change the channel according to the dataset.
LEVELS = 10
channel = 1
model = Net(in_channel1=channel * LEVELS,
out_channel1=32 * LEVELS,
out_channel2=64 * LEVELS).to(device)
optimizer = optim.SGD(model.parameters(), lr=0.0001, momentum=0.2)
logger.info('Load model.')
save_model = True
for epoch in range(1, 50 + 1): ## 5 batches
print('Running epoch ', epoch)
train(model, device, train_loader, optimizer, epoch)
test(model, device, test_loader)
if (save_model):
torch.save(model.state_dict(),
"deeprobust/image/save_models/thermometer_encoding.pt")
