def run():
print("CUDA is available: {}".format(torch.cuda.is_available()))
data_transform = transforms.Compose(
[Rescale(250), CenterCrop(224),
Normalize(), ToTensor()])
# loader will split datatests into batches witht size defined by batch_size
train_loader = initialize_train_loader(data_transform)
test_loader = initialize_test_loader(data_transform)
model_id = time.strftime("%Y-%m-%dT%H:%M:%S", time.localtime())
# instantiate the neural network
net = Net()
net.to(device=device)
summary(net, (1, 224, 224))
# define the loss function using SmoothL1Loss
criterion = nn.SmoothL1Loss()
# define the params updating function using Adam
optimizer = optim.Adam(net.parameters(), lr=0.001)
loss_logger = []
for i in range(1, epochs + 1):
model_name = 'model-{}-epoch-{}.pt'.format(model_id, i)
# train all data for one epoch
train(net, criterion, optimizer, i, train_loader, model_id,
loss_logger)
# evaludate the accuracy after each epoch
evaluate(net, criterion, i, test_loader)
# save model after every 5 epochs
# https://discuss.pytorch.org/t/loading-a-saved-model-for-continue-training/17244/3
# https://github.com/pytorch/pytorch/issues/2830
# https://pytorch.org/tutorials/beginner/saving_loading_models.html
if i % 5 == 1:
torch.save(
{
'epoch': i,
'model': net.state_dict(),
'optimizer': optimizer.state_dict(),
'loss_logger': loss_logger,
}, model_dir + model_name)
print("Finished training!")
model_name = 'model-{}-final.pt'.format(model_id)
torch.save(
{
'epoch': epochs,
'model': net.state_dict(),
'optimizer': optimizer.state_dict(),
'loss_logger': loss_logger,
}, model_dir + model_name)
def main(args):
train_loader, val_loader = create_dataloaders(args.batch_size)
model = Net().to(device)
optim = torch.optim.Adam(model.parameters())
lr_schedule = torch.optim.lr_scheduler.ReduceLROnPlateau(optim,
patience=1,
verbose=True)
criterion = torch.nn.CrossEntropyLoss()
best_accuracy = 0
for epoch in range(1, args.epochs + 1):
model.train()
train_loss, train_accuracy = do_epoch(model,
train_loader,
criterion,
optim=optim)
model.eval()
with torch.no_grad():
val_loss, val_accuracy = do_epoch(model,
val_loader,
criterion,
optim=None)
tqdm.write(
f'EPOCH {epoch:03d}: train_loss={train_loss:.4f}, train_accuracy={train_accuracy:.4f} '
f'val_loss={val_loss:.4f}, val_accuracy={val_accuracy:.4f}')
if val_accuracy > best_accuracy:
print('Saving model...')
best_accuracy = val_accuracy
torch.save(model.state_dict(), 'trained_models/source.pt')
lr_schedule.step(val_loss)
def main():
# Training settings
parser = argparse.ArgumentParser(description='PyTorch Vec2Color Example')
parser.add_argument('--batch-size', type=int, default=64, metavar='N',
help='input batch size for training (default: 64)')
parser.add_argument('--test-batch-size', type=int, default=1000, metavar='N',
help='input batch size for testing (default: 1000)')
parser.add_argument('--epochs', type=int, default=1500, metavar='N',
help='number of epochs to train (default: 1500)')
parser.add_argument('--lr', type=float, default=1.0, metavar='LR',
help='learning rate (default: 1.0)')
parser.add_argument('--gamma', type=float, default=0.7, metavar='M',
help='Learning rate step gamma (default: 0.7)')
parser.add_argument('--no-cuda', action='store_true', default=False,
help='disables CUDA training')
parser.add_argument('--seed', type=int, default=1, metavar='S',
help='random seed (default: 1)')
parser.add_argument('--log-interval', type=int, default=10, metavar='N',
help='how many batches to wait before logging training status')
parser.add_argument('--save-model', action='store_true', default=False,
help='For Saving the current Model')
args = parser.parse_args()
use_cuda = not args.no_cuda and torch.cuda.is_available()
torch.manual_seed(args.seed)
device = torch.device("cuda" if use_cuda else "cpu")
kwargs = {'num_workers': 1, 'pin_memory': True} if use_cuda else {}
file_names = ('capitalize', 'lower', 'upper', 'title')
x_df = pd.concat([pd.read_csv('doc2color/data/{}.csv'.format(file_name)) for file_name in file_names])
y_df = pd.concat([pd.read_csv('doc2color/data/rgb.csv')] * len(file_names))
tensor_x = torch.stack([torch.from_numpy(np.array(i)) for i in x_df.values.astype(np.float32)])
tensor_y = torch.stack([torch.from_numpy(np.array(i)) for i in y_df.values.astype(np.float32) / 255.0])
x_train, x_test, y_train, y_test = train_test_split(
tensor_x, tensor_y, test_size=0.01, random_state=args.seed)
train_dataset = torch.utils.data.TensorDataset(x_train, y_train)
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=args.batch_size, shuffle=True, **kwargs)
test_dataset = torch.utils.data.TensorDataset(x_test, y_test)
test_loader = torch.utils.data.DataLoader(test_dataset,
batch_size=args.test_batch_size, shuffle=True, **kwargs)
model = Net().to(device)
optimizer = optim.Adadelta(model.parameters(), lr=args.lr)
scheduler = StepLR(optimizer, step_size=1, gamma=args.gamma)
for epoch in range(1, args.epochs + 1):
train(args, model, device, train_loader, optimizer, epoch)
test(args, model, device, test_loader)
scheduler.step()
if args.save_model:
torch.save(model.state_dict(), "doc2color/pt_objects/vec2color.pt")
def main():
# Training settings
parser = argparse.ArgumentParser(description="PyTorch MNIST Example")
parser.add_argument(
"--batch-size",
type=int,
default=1,
metavar="N",
help="input batch size for training (default: 1)",
)
parser.add_argument(
"--test-batch-size",
type=int,
default=1,
metavar="N",
help="input batch size for testing (default: 1)",
)
parser.add_argument(
"--epochs",
type=int,
default=3,
metavar="N",
help="number of epochs to train (default: 14)",
)
parser.add_argument(
"--lr",
type=float,
default=1.0,
metavar="LR",
help="learning rate (default: 1.0)",
)
parser.add_argument(
"--gamma",
type=float,
default=0.7,
metavar="M",
help="Learning rate step gamma (default: 0.7)",
)
parser.add_argument("--no-cuda",
action="store_true",
default=False,
help="disables CUDA training")
parser.add_argument(
"--dry-run",
action="store_true",
default=False,
help="quickly check a single pass",
)
parser.add_argument("--seed",
type=int,
default=1,
metavar="S",
help="random seed (default: 1)")
parser.add_argument(
"--log-interval",
type=int,
default=100,
metavar="N",
help="how many batches to wait before logging training status",
)
parser.add_argument(
"--save-model",
action="store_true",
default=False,
help="For Saving the current Model",
)
args = parser.parse_args()
use_cuda = not args.no_cuda and torch.cuda.is_available()
torch.manual_seed(args.seed)
device = torch.device("cuda" if use_cuda else "cpu")
kwargs = {"batch_size": args.batch_size}
if use_cuda:
kwargs.update({
"num_workers": 1,
"pin_memory": True,
"shuffle": True
}, )
transform = transforms.Compose(
[transforms.ToTensor(),
transforms.Normalize((0.1307, ), (0.3081, ))])
dataset1 = datasets.MNIST("../data",
train=True,
download=True,
transform=transform)
dataset2 = datasets.MNIST("../data", train=False, transform=transform)
train_loader = torch.utils.data.DataLoader(dataset1, **kwargs)
test_loader = torch.utils.data.DataLoader(dataset2, **kwargs)
model = Net().to(device)
for epoch in range(1, args.epochs + 1):
train(args, model, device, train_loader, epoch)
test(model, device, test_loader)
if args.save_model:
torch.save(model.state_dict(), "mnist_cnn.pt")
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--datapath",
type=str,
default="data/3droad.mat",
help="dataset file")
parser.add_argument('--epochs',
type=int,
default=30,
help='number of epochs to train')
parser.add_argument('--batch_size',
type=int,
default=1024,
help='input batch size for training')
parser.add_argument('--val_batch_size',
type=int,
default=1024,
help='input batch size for valing')
# Training Settings
parser.add_argument('--arch_file',
type=str,
default=None,
help='name of file which defines the architecture')
parser.add_argument('--arch_name',
type=str,
default='LeNet5',
help='name of the architecture')
parser.add_argument('--arch_args',
type=json.loads,
default=None,
help='[JSON] arguments for the architecture')
parser.add_argument('--optim_name',
type=str,
default=VIOptimizer.__name__,
help='name of the optimizer')
parser.add_argument('--optim_args',
type=json.loads,
default=None,
help='[JSON] arguments for the optimizer')
parser.add_argument('--curv_args',
type=json.loads,
default=dict(),
help='[JSON] arguments for the curvature')
parser.add_argument('--fisher_args',
type=json.loads,
default=dict(),
help='[JSON] arguments for the fisher')
parser.add_argument('--scheduler_name',
type=str,
default=None,
help='name of the learning rate scheduler')
parser.add_argument('--scheduler_args',
type=json.loads,
default=None,
help='[JSON] arguments for the scheduler')
# Options
parser.add_argument('--no_cuda',
action='store_true',
default=False,
help='disables CUDA training')
parser.add_argument('--seed', type=int, default=1, help='random seed')
parser.add_argument(
'--log_interval',
type=int,
default=20,
help='how many batches to wait before logging training status')
parser.add_argument(
'--checkpoint_interval',
type=int,
default=50,
help='how many epochs to wait before logging training status')
parser.add_argument('--out',
type=str,
default='results/regression/',
help='dir to save output files')
parser.add_argument('--config', default=None, help='config file path')
parser.add_argument("--log_name",
default=None,
required=True,
type=str,
help="log name")
args = parser.parse_args()
dict_args = vars(args)
use_cuda = not args.no_cuda and torch.cuda.is_available()
device = torch.device("cuda" if use_cuda else "cpu")
if args.config is not None:
with open(args.config) as f:
config = json.load(f)
dict_args.update(config)
if not os.path.isfile(args.datapath):
print('Downloading \'3droad\' UCI dataset...')
urllib.request.urlretrieve(
'https://www.dropbox.com/s/f6ow1i59oqx05pl/3droad.mat?dl=1',
args.datapath)
data = loadmat(args.datapath)['data']
X = data[:, :-1]
X = X - X.min(0)[0]
X = 2 * (X / X.max(0)[0]) - 1
y = data[:, -1]
train_val_n = int(floor(0.8 * len(X)))
X_train_val = X[:train_val_n]
y_train_val = y[:train_val_n]
X_test = X[train_val_n:, :]
y_test = y[train_val_n:]
X_train, X_val, y_train, y_val = train_test_split(X_train_val,
y_train_val,
test_size=0.2)
dtype = torch.float32
X_train = torch.tensor(X_train, dtype=dtype, requires_grad=False)
y_train = torch.tensor(y_train, dtype=dtype, requires_grad=False)
X_val = torch.tensor(X_val, dtype=dtype, requires_grad=False)
y_val = torch.tensor(y_val, dtype=dtype, requires_grad=False)
X_test = torch.tensor(X[train_val_n:, :], dtype=dtype, requires_grad=False)
y_test = torch.tensor(y[train_val_n:], dtype=dtype, requires_grad=False)
train_dataset = TensorDataset(X_train, y_train)
train_loader = DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=True)
val_dataset = TensorDataset(X_val, y_val)
val_loader = DataLoader(val_dataset,
batch_size=args.val_batch_size,
shuffle=True)
test_dataset = TensorDataset(X_test, y_test)
test_loader = DataLoader(test_dataset,
batch_size=args.val_batch_size,
shuffle=True)
data_dim = X_train.size(-1)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = Net(data_dim).to(device).float()
optim_kwargs = {} if args.optim_args is None else args.optim_args
if args.optim_name == SecondOrderOptimizer.__name__:
optimizer = SecondOrderOptimizer(model,
**config["optim_args"],
curv_kwargs=config["curv_args"])
elif args.optim_name == VIOptimizer.__name__:
optimizer = VIOptimizer(model,
dataset_size=len(train_loader.dataset),
seed=args.seed,
**config["optim_args"],
curv_kwargs=config["curv_args"])
else:
modules = import_module("optimizers")
optim_class = getattr(modules, args.optim_name)
optimizer = optim_class(model.parameters())
if args.scheduler_name is None:
scheduler = None
else:
scheduler_class = getattr(torchsso.optim.lr_scheduler,
args.scheduler_name, None)
if scheduler_class is None:
scheduler_class = getattr(torch.optim.lr_scheduler,
args.scheduler_name)
scheduler_kwargs = config["scheduler_args"]
scheduler = scheduler_class(optimizer, **scheduler_kwargs)
log_file_name = "log_" + args.log_name
logger = Logger(args.out, log_file_name)
logger.start()
# train
epochs = 120
model.train()
print("=========== Start ===========")
for i in range(args.epochs):
losses = 0
for minibatch_i, (x_batch, y_batch) in enumerate(train_loader):
x_batch = x_batch.to(device)
y_batch = y_batch.to(device)
"""
def closure():
optimizer.zero_grad()
output = model(x_batch)
loss = F.mse_loss(output, y_batch, reduction="sum").float()
loss.backward()
return loss, output
"""
optimizer.zero_grad()
output = model(x_batch)
loss = F.mse_loss(output, y_batch, reduction="sum").float()
loss.backward()
optimizer.step()
losses += loss.item()
if (minibatch_i + 1) % args.log_interval == 0:
print(
"Train Epoch: {} batch idx: {} elapsed time: {:.1f}s MSE: {}"
.format(i + 1, minibatch_i + 1, logger.elapsed_time,
loss.item() / args.batch_size))
losses = losses / len(train_loader.dataset)
val_mse = validate(model, device, val_loader, optimizer, "Eval")
iteration = (i + 1) * len(train_loader)
log = {
"epoch": i + 1,
"iteration": iteration,
"mse": losses,
"val_mse": val_mse,
"lr": optimizer.param_groups[0]["lr"],
"momentum": optimizer.param_groups[0].get("momentum", 0)
}
logger.write(log)
if i % args.checkpoint_interval == 0 or i + 1 == args.epochs:
path = os.path.join(args.out,
"epoch{}_{}.ckpt".format(i + 1, args.log_name))
data = {
"model": model.state_dict(),
"optimizer": optimizer.state_dict(),
"epoch": i + 1
}
torch.save(data, path)
print("=========== Test ===========")
test_mse = validate(model, device, test_loader, optimizer, "Test")
log = {"test_mse": test_mse}
logger.write(log)
class TrainBigramNN(tune.Trainable):
def _setup(self, config):
print("Loading word vectors...")
word2index, word_vecs = process_word_vecs(FAST_TEXT)
# Note that the word embeddings are normalized.
self.wv = WV(F.normalize(word_vecs), word2index)
# wv = WV(word_vecs, word2index)
print("Done.")
self.corpus_size = config["corpus_size"]
bigram_fn_name = "diff"
out_bigram_dim = 300
dist_fn_name = "cos_dist"
loss_fn_name = "mrl"
margin = config["margin"]
self.lr = config["lr"]
self.num_epochs = config["num_epochs"]
self.batch_size = config["batch_size"]
self.test_model = True
self.test_freq = config["test_freq"]
with open(PROCESSED / "train.{}.pkl".format(str(self.corpus_size)), "rb") as f:
wiki_train = pickle.load(f)
with open(PROCESSED / "valid.pkl", "rb") as f:
wiki_valid = pickle.load(f)
wiki_combined = wiki_train + wiki_valid
self.corpus = Corpus("wiki", wiki_combined, self.wv, filter_stopwords=True)
self.model = Net(
self.wv.vecs.size(1), BigramEncoder(bigram_fn_name), out_bigram_dim
)
self.model.to(device)
self.dist_fn = DistanceFunction(dist_fn_name)
self.loss_fn = LossFunction(loss_fn_name, margin=margin)
self.device = device
self.optimizer = torch.optim.Adam(self.model.parameters(), lr=self.lr)
torch.manual_seed(config["seed"])
print("Traninig on Wikipedia corpus of size {}".format(self.corpus_size))
def _train(self):
result = train(
self.wv,
self.corpus.ix_sents[: self.corpus_size],
self.corpus.sent_lengths[: self.corpus_size],
self.corpus.ix_sents[self.corpus_size :],
self.corpus.sent_lengths[self.corpus_size :],
self.model,
self.wv.vecs,
self.dist_fn,
self.loss_fn,
self.optimizer,
self.lr,
self.num_epochs,
self.batch_size,
self._iteration,
self.test_model,
self.test_freq,
self.device,
)
return result
def _save(self, tmp_checkpoint_dir):
checkpoint_path = str(Path(tmp_checkpoint_dir) / "model.pth")
torch.save(self.model.state_dict(), checkpoint_path)
return checkpoint_path
def _restore(self, tmp_checkpoint_dir):
checkpoint_path = str(Path(tmp_checkpoint_dir) / "model.pth")
self.model.load_state_dict(torch.load(checkpoint_path))
output_pts = net(images)
loss = criterion(output_pts, key_pts)
optimizer.zero_grad()
loss.backward()
optimizer.step()
train_loss += loss.item()
with torch.no_grad():
for data in test_loader:
images = data["image"]
key_pts = data["keypoints"]
key_pts = key_pts.view(key_pts.size(0), -1)
key_pts = key_pts.type(dtype)
images = images.type(dtype)
output_pts = net(images)
loss = criterion(output_pts, key_pts)
test_loss += loss.item()
print('Epoch: {}, trainloss: {}, testloss {}'.format(
epoch + 1, train_loss / len(train_loader),
test_loss / len(test_loader)))
print('Finished Training')
train_net(42)
torch.save(net.state_dict(), "v2_model.pt")
def test(epoch):
model.eval() # 把module设置为评估模式，只对Dropout和BatchNorm模块有影响
test_loss = 0
correct = 0
for data, target in test_loader:
if args.cuda:
data, target = data.cuda(), target.cuda()
data, target = Variable(data, volatile=True), Variable(target)
output = model(data)
test_loss += F.nll_loss(output, target).data[0] # Variable.data
# get the index of the max log-probability
pred = output.data.max(1)[1]
correct += pred.eq(target.data).cpu().sum()
test_loss = test_loss
# loss function already averages over batch size
test_loss /= len(test_loader)
print(
'\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
test_loss, correct, len(test_loader.dataset),
100. * correct / len(test_loader.dataset)))
for epoch in range(1, args.epochs + 1):
train(epoch)
test(epoch)
if ((epoch + 1) % args.save_interval) == 0:
path_checkpoint = '{0}/{1}_state_epoch{2}.pth'.format(
'checkpoints', 'mnist', epoch + 1)
utils.save_checkpoint(model.state_dict(), path_checkpoint)
def main():
# Define network
net = Net()
print(net)
data_transform = transforms.Compose(
[Rescale(250), Normalize(), RandomCrop(224)])
train_dataset = FacialKeypointsDataset(
csv_file='data/training_frames_keypoints.csv',
root_dir='data/training/',
transform=data_transform)
# iterate through the transformed dataset and print some stats about the first few samples
for i in range(4):
sample = train_dataset[i]
print(i, sample['image'].size, sample['keypoints'].size)
train_loader = DataLoader(train_dataset,
batch_size=10,
shuffle=True,
num_workers=4)
test_dataset = FacialKeypointsDataset(
csv_file='data/test_frames_keypoints.csv',
root_dir='data/test/',
transform=data_transform)
test_loader = DataLoader(test_dataset,
batch_size=1,
shuffle=True,
num_workers=4)
test_images, test_outputs, gt_pts = net_sample_output(net, test_loader)
# print out the dimensions of the data to see if they make sense
print(test_images.data.size())
print(test_outputs.data.size())
print(gt_pts.size())
# call it
visualize_output(test_images, test_outputs, gt_pts, 1)
criterion = nn.MSELoss()
optimizer = optim.Adam(net.parameters(),
lr=0.001,
betas=(0.9, 0.999),
eps=1e-8)
n_epochs = 2
train_net(net, criterion, optimizer, train_loader, n_epochs)
# get a sample of test data again
test_images, test_outputs, gt_pts = net_sample_output(net, test_loader)
print(test_images.data.size())
print(test_outputs.data.size())
print(gt_pts.size())
model_dir = 'saved_models/'
model_name = 'keypoints_model_1.pt'
# after training, save your model parameters in the dir 'saved_models'
torch.save(net.state_dict(), model_dir + model_name)
weights1 = net.conv1.weight.data
w = weights1.numpy()
filter_index = 0
print(w[filter_index][0])
print(w[filter_index][0].shape)
# display the filter weights
plt.imshow(w[filter_index][0], cmap='gray')
# Make sure the relevant folders have been created -------------
if not os.path.exists('disturbance-fig'):
os.makedirs('disturbance-fig')
if not os.path.exists('weights'):
os.makedirs('weights')
# -----------------------------------------------------------------
# name to be used for saving the network weights and perturbation data
out_filename = 'CE_with_sd_coef_of_' + str(args.sd_coef)
if not args.resume:
fout_weights = os.path.join('weights', out_filename + '.pt')
torch.save(model.state_dict(), fout_weights)
print('Written to: {}'.format(fout_weights))
model.eval()
model.cuda()
foolbox_model = foolbox.models.PyTorchModel(model=model,
num_classes=2,
bounds=(-1, 1))
criterion = foolbox.criteria.Misclassification()
attack = foolbox.attacks.PGD(foolbox_model,
criterion,
distance=foolbox.distances.Linfinity)
def get_adv_disturbance_mag(attack, dataset, true_label_for_dataset, order):
list_of_mag = []
#inputs, labels = Variable(inputs), Variable(labels)
#print( inputs.size() )
#outputs = model(inputs)
#print( outputs.size() )
#
#
#import sys
#sys.exit()
criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
scheduler = lr_scheduler.StepLR(optimizer, step_size=7, gamma=0.1)
dloaders = {'train': train_loader, 'valid': valid_loader}
best_model_wts = model.state_dict()
best_acc = 0.0
dataset_sizes = {
'train': len(dloaders['train'].dataset),
'valid': len(dloaders['valid'].dataset)
}
for epoch in range(num_epochs):
print('Epoch', epoch)
for phase in ['train', 'valid']:
if phase == 'train':
scheduler.step()
model.train(True)
else:
model.train(False)
running_loss = 0.0
def main(args):
clf_model = Net().to(device)
clf_model.load_state_dict(torch.load(args.MODEL_FILE))
feature_extractor = clf_model.feature_extractor
discriminator = clf_model.classifier
critic = nn.Sequential(nn.Linear(320, 50), nn.ReLU(), nn.Linear(50, 20),
nn.ReLU(), nn.Linear(20, 1)).to(device)
half_batch = args.batch_size // 2
if args.adapt_setting == 'mnist2mnistm':
source_dataset = MNIST(config.DATA_DIR / 'mnist',
train=True,
download=True,
transform=Compose(
[GrayscaleToRgb(),
ToTensor()]))
target_dataset = MNISTM(train=False)
elif args.adapt_setting == 'svhn2mnist':
source_dataset = ImageClassdata(txt_file=args.src_list,
root_dir=args.src_root,
img_type=args.img_type,
transform=transforms.Compose([
transforms.Resize(28),
transforms.ToTensor(),
]))
target_dataset = ImageClassdata(txt_file=args.tar_list,
root_dir=args.tar_root,
img_type=args.img_type,
transform=transforms.Compose([
transforms.ToTensor(),
]))
elif args.adapt_setting == 'mnist2usps':
source_dataset = ImageClassdata(txt_file=args.src_list,
root_dir=args.src_root,
img_type=args.img_type,
transform=transforms.Compose([
transforms.ToTensor(),
]))
target_dataset = ImageClassdata(txt_file=args.tar_list,
root_dir=args.tar_root,
img_type=args.img_type,
transform=transforms.Compose([
transforms.Resize(28),
transforms.ToTensor(),
]))
else:
raise NotImplementedError
source_loader = DataLoader(source_dataset,
batch_size=half_batch,
shuffle=True,
num_workers=1,
pin_memory=True,
drop_last=True)
target_loader = DataLoader(target_dataset,
batch_size=half_batch,
shuffle=True,
num_workers=1,
pin_memory=True,
drop_last=True)
critic_optim = torch.optim.Adam(critic.parameters(), lr=1e-4)
clf_optim = torch.optim.Adam(clf_model.parameters(), lr=1e-4)
clf_criterion = nn.CrossEntropyLoss()
if not os.path.exists('logs'): os.makedirs('logs')
f = open(f'logs/{args.adapt_setting}_{args.name}.txt', 'w+')
for epoch in range(1, args.epochs + 1):
batch_iterator = zip(loop_iterable(source_loader),
loop_iterable(target_loader))
total_loss = 0
total_accuracy = 0
target_label_accuracy = 0
for _ in trange(args.iterations, leave=False):
(source_x, source_y), (target_x, target_y) = next(batch_iterator)
# Train critic
set_requires_grad(feature_extractor, requires_grad=False)
set_requires_grad(critic, requires_grad=True)
source_x, target_x = source_x.to(device), target_x.to(device)
source_y = source_y.to(device)
with torch.no_grad():
h_s = feature_extractor(source_x).data.view(
source_x.shape[0], -1)
h_t = feature_extractor(target_x).data.view(
target_x.shape[0], -1)
for _ in range(args.k_critic):
gp = gradient_penalty(critic, h_s, h_t)
critic_s = critic(h_s)
critic_t = critic(h_t)
wasserstein_distance = critic_s.mean() - critic_t.mean()
critic_cost = -wasserstein_distance + args.gamma * gp
critic_optim.zero_grad()
critic_cost.backward()
critic_optim.step()
total_loss += critic_cost.item()
# Train classifier
set_requires_grad(feature_extractor, requires_grad=True)
set_requires_grad(critic, requires_grad=False)
for _ in range(args.k_clf):
source_features = feature_extractor(source_x).view(
source_x.shape[0], -1)
target_features = feature_extractor(target_x).view(
target_x.shape[0], -1)
source_preds = discriminator(source_features)
clf_loss = clf_criterion(source_preds, source_y)
wasserstein_distance = critic(source_features).mean() - critic(
target_features).mean()
loss = clf_loss + args.wd_clf * wasserstein_distance
clf_optim.zero_grad()
loss.backward()
clf_optim.step()
target_preds = discriminator(target_features)
target_label_accuracy += (target_preds.cpu().max(1)[1] ==
target_y).float().mean().item()
mean_loss = total_loss / (args.iterations * args.k_critic)
# mean_accuracy = total_accuracy / (args.iterations * args.k_disc)
target_mean_accuracy = target_label_accuracy / (args.iterations *
args.k_clf)
tqdm.write(
f'EPOCH {epoch:03d}: critic_loss={mean_loss:.4f}, target_accuracy={target_mean_accuracy:.4f}'
)
f.write(
f'EPOCH {epoch:03d}: critic_loss={mean_loss:.4f}, target_accuracy={target_mean_accuracy:.4f}'
)
torch.save(
clf_model.state_dict(),
f'trained_models/{args.adapt_setting}_{args.name}_ep{epoch}.pt')
f.close()
if __name__ == '__main__':
if MODE == "train":
print('Size of training datset: ', trainset_length)
if not os.path.exists(os.path.join(os.getcwd(), exp_name, 'weights')):
os.makedirs(os.path.join(os.getcwd(), exp_name, 'weights'))
print('Training Started!')
for e in range(NUM_EPOCHS):
train(e)
r_scheduler.step(e)
d_scheduler.step(e)
if e % 2 == 0:
torch.save(
refine_net.state_dict(),
os.path.join(os.getcwd(), exp_name, 'weights',
"weights_net_epoch_{}.pth".format(e)))
torch.save(
d_net.state_dict(),
os.path.join(os.getcwd(), exp_name, 'weights',
"weights_net_epoch_disc_{}.pth".format(e)))
if e == 0:
cmd = f"gsutil -m cp -r {exp_name}/ {save_path}/"
os.system(cmd)
else:
try:
cmd = f"gsutil -m cp -r {exp_name}/weights/weights_net_epoch_{e}.pth {save_path}/{exp_name}/weights/"
os.system(cmd)
cmd = f"gsutil -m cp -r {exp_name}/weights/weights_net_epoch_disc_{e}.pth {save_path}/{exp_name}/weights/"
os.system(cmd)
loss = criterion(output_pts, key_pts)
# zero the parameter (weight) gradients
optimizer.zero_grad()
# backward pass to calculate the weight gradients
loss.backward()
# update the weights
optimizer.step()
# print loss statistics
running_loss += loss.item()
if batch_i % 10 == 9: # print every 10 batches
print('Epoch: {}, Batch: {}, Avg. Loss: {}'.format(epoch + 1, batch_i + 1, running_loss / 1000))
running_loss = 0.0
print('Finished Training')
# train your network
n_epochs = 100 # start small, and increase when you've decided on your model structure and hyperparams
# this is a Workspaces-specific context manager to keep the connection
# alive while training your model, not part of pytorch
#with active_session():
train_net(n_epochs)
PATH = r'C:\Users\Semanti Basu\Documents\OneDrive_2020-02-19\3D Ceaser dataset\Image and point generation\Image and point generation\frontaltrainedmodel_10epoch.pth'
torch.save(net.state_dict(), PATH)
def train(options):
exp_name = options['exp_name']
batch_size = options['batch_size']
use_pca = options['use_pca']
model_type = options['model_type']
loss_fn = options['loss_fn']
optim = options['optim']
use_scheduler = options['use_scheduler']
lr = options['lr']
epochs = options['epochs']
pca_var_hold = options['pca_var_hold']
debug_mode = options['debug_mode']
if os.path.exists(exp_name):
shutil.rmtree(exp_name)
time.sleep(1)
writer = SummaryWriter(exp_name,flush_secs=1)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
X = os.listdir('hilbert_data')
X_train = X[:int(0.8*len(X))]
X_test = X[int(0.8*len(X)):]
# X = np.load('bined_x.npy')
# y = np.load('bined_y.npy')
# X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
# if use_pca and 'Raw' in exp_name:
# scaler = PCA(pca_var_hold)
# scaler.fit(X_train)
# X_train = scaler.transform(X_train)
# X_test = scaler.transform(X_test)
# needed_dim = X_train.shape[1]
dataset_train = HIL_MOOD(X_train, model_type=model_type,data_type='train',debug_mode=debug_mode)
train_loader = DataLoader(dataset=dataset_train, batch_size=batch_size, shuffle=True)
dataset_val = HIL_MOOD(X_test, model_type=model_type,data_type='val')
valid_loader = DataLoader(dataset=dataset_val, batch_size=batch_size, shuffle=False)
model = Net()
model.to(device)
if optim == None:
print('you need to specify an optimizer')
exit()
elif optim == 'adam':
optimizer = torch.optim.Adam( model.parameters(), lr=lr)
elif optim == 'sgd':
optimizer = torch.optim.SGD( model.parameters(), lr=lr,momentum=0.9)
if use_scheduler:
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min',verbose=True,threshold=0.0001,patience = 10)
if loss_fn == None:
print('you need to specify an optimizer')
exit()
else:
if loss_fn == 'mse':
loss_fn = torch.nn.MSELoss()
elif loss_fn == 'cross_entropy':
loss_fn = torch.nn.CrossEntropyLoss()
mean_train_losses = []
mean_valid_losses = []
valid_acc_list = []
best = 0 #small number for acc big number for loss to save a model
for epoch in range(epochs):
model.train()
train_losses = []
valid_losses = []
for i, (images, labels) in enumerate(train_loader):
if images.shape[0] != batch_size:
continue
images = images.to(device)
labels = labels.to(device)
optimizer.zero_grad()
# print(images.shape)
outputs = model(images)
# print(images.shape)
# print(outputs.shape)
# print(labels.shape)
# print(i)
loss =loss_fn(outputs,labels)
# print('loss: ',loss.item())
writer.add_scalar('Loss/train', loss.item(), len(train_loader)*epoch+i)
loss.backward()
optimizer.step()
train_losses.append(loss.item())
del outputs
# if (i * batch_size) % (batch_size * 100) == 0:
# print(f'{i * batch_size} / 50000')
model.eval()
correct_5_2 = 0
correct_5_1 = 0
total_loss = 0
total = 0
accsat =[0.5,0.05,0.005]
accs = np.zeros(len(accsat))
# corrs = np.zeros(len(accsat))
correct_array = np.zeros(len(accsat))
with torch.no_grad():
for i, (images, labels) in enumerate(valid_loader):
images = images.to(device)
labels = labels.to(device)
outputs = model(images)
loss = loss_fn(outputs, labels)
for i in range(len(accsat)):
correct_array[i] += accat(outputs,labels,thresh=accsat[i])
# total_loss += loss.item()
total += labels.size(0)
valid_losses.append(loss.item())
mean_train_losses.append(np.mean(train_losses))
mean_valid_losses.append(np.mean(valid_losses))
# scheduler.step(np.mean(valid_losses))
for i in range(len(accsat)):
accs[i] = 100*correct_array[i]/total
writer.add_scalar('Acc/val_@'+str(accsat[i]), accs[i], epoch)
if np.mean(valid_losses) < best:
best = np.mean(valid_losses)
torch.save(model.state_dict(),os.path.join(os.getcwd(),'models','meh.pth'))
writer.add_scalar('Loss/val', np.mean(valid_losses), epoch)
# valid_acc_list.append(accuracy)
if epoch ==epochs-1:
print('epoch : {}, train loss : {:.4f}, valid loss : {:.4f}, [email protected] : {:.4f}'\
.format(epoch+1, np.mean(train_losses), np.mean(valid_losses), accsat[1]))
def main():
check_dir = '../LPSfiles/' + name
if not os.path.exists(check_dir):
os.mkdir(check_dir)
# data
val_loader = torch.utils.data.DataLoader(PriorFolder(opt.val_dir,
opt.prior_dir,
size=256,
mean=mean,
std=std),
batch_size=opt.b * 3,
shuffle=False,
num_workers=4,
pin_memory=True)
train_loader = torch.utils.data.DataLoader(Folder(opt.train_dir,
scales=[64] * 3 +
[128, 256],
crop=0.9,
flip=True,
rotate=None,
mean=mean,
std=std),
collate_fn=collate_more,
batch_size=opt.b * 6,
shuffle=True,
num_workers=4,
pin_memory=True)
# models
p = 5
net = Net(base=opt.base)
fcn = FCN(net)
net = nn.DataParallel(net).cuda()
net.train()
"""
# fcn = nn.DataParallel(fcn).cuda()
# sdict =torch.load('/home/crow/LPSfiles/Train2_vgg16/fcn-iter13800.pth')
# fcn.load_state_dict(sdict)
fcn = nn.DataParallel(fcn).cuda()
fcn.train()
optimizer = torch.optim.Adam([
{'params': fcn.parameters(), 'lr': 1e-4},
])
logs = {'best_it':0, 'best': 0}
sal_data_iter = iter(train_loader)
i_sal_data = 0
for it in tqdm(range(opt.max)):
# for it in tqdm(range(1)):
# if it > 1000 and it % 100 == 0:
# optimizer.param_groups[0]['lr'] *= 0.5
if i_sal_data >= len(train_loader):
sal_data_iter = iter(train_loader)
i_sal_data = 0
data, lbls, _ = sal_data_iter.next()
i_sal_data += 1
data = data.cuda()
lbls = [lbl.unsqueeze(1).cuda() for lbl in lbls]
msks = fcn(data)
loss = sum([F.binary_cross_entropy_with_logits(msk, lbl) for msk, lbl in zip(msks, lbls)])
optimizer.zero_grad()
loss.backward()
optimizer.step()
if it % 10 == 0:
writer.add_scalar('loss', loss.item(), it)
image = make_image_grid(data[:6], mean, std)
writer.add_image('Image', torchvision.utils.make_grid(image), it)
big_msk = F.sigmoid(msks[-1]).expand(-1, 3, -1, -1)
writer.add_image('msk', torchvision.utils.make_grid(big_msk.data[:6]), it)
big_msk = lbls[-1].expand(-1, 3, -1, -1)
writer.add_image('gt', torchvision.utils.make_grid(big_msk.data[:6]), it)
# if it % 100 == 0:
if it != 0 and it % 100 == 0:
fm, mae = validate(val_loader, fcn, os.path.join(check_dir, 'results'),
os.path.join(opt.val_dir, 'masks'))
print(u'损失: %.4f'%(loss.item()))
print(u'最大FM: iteration %d的%.4f, 这次FM: %.4f'%(logs['best_it'], logs['best'], fm))
logs[it] = {'FM': fm}
if fm > logs['best']:
logs['best'] = fm
logs['best_it'] = it
torch.save(fcn.state_dict(), '%s/fcn-best.pth' % (check_dir))
with open(os.path.join(check_dir, 'logs.json'), 'w') as outfile:
json.dump(logs, outfile)
torch.save(fcn.state_dict(), '%s/fcn-iter%d.pth' % (check_dir, it))
"""
###################################################################################################
val_loader = torch.utils.data.DataLoader(PriorFolder(opt.val_dir,
opt.prior_dir,
size=256,
mean=mean,
std=std),
batch_size=opt.b,
shuffle=False,
num_workers=4,
pin_memory=True)
train_loader = torch.utils.data.DataLoader(Folder(opt.train_dir,
scales=[256],
crop=0.9,
flip=True,
rotate=None,
mean=mean,
std=std),
collate_fn=collate_more,
batch_size=opt.b,
shuffle=True,
num_workers=4,
pin_memory=True)
optimizer = torch.optim.Adam([
{
'params': net.parameters(),
'lr': 1e-4
},
])
logs = {'best_it': 0, 'best': 0}
sal_data_iter = iter(train_loader)
i_sal_data = 0
for it in tqdm(range(opt.max)):
# if it > 1000 and it % 100 == 0:
# optimizer.param_groups[0]['lr'] *= 0.5
if i_sal_data >= len(train_loader):
sal_data_iter = iter(train_loader)
i_sal_data = 0
data, lbl, _ = sal_data_iter.next()
i_sal_data += 1
data = data.cuda()
lbl = lbl[0].unsqueeze(1)
noisy_label = (lbl.numpy() +
np.random.binomial(1,
float(p) / 100.0, (256, 256))) % 2
noisy_label = torch.Tensor(noisy_label).cuda()
lbl = lbl.cuda()
msk = net(data, noisy_label)
loss = F.binary_cross_entropy_with_logits(msk, lbl)
optimizer.zero_grad()
loss.backward()
optimizer.step()
if it % 10 == 0:
writer.add_scalar('loss', loss.item(), it)
image = make_image_grid(data[:6], mean, std)
writer.add_image('Image', torchvision.utils.make_grid(image), it)
big_msk = F.sigmoid(msk).expand(-1, 3, -1, -1)
writer.add_image('msk',
torchvision.utils.make_grid(big_msk.data[:6]), it)
big_msk = lbl.expand(-1, 3, -1, -1)
writer.add_image('gt',
torchvision.utils.make_grid(big_msk.data[:6]), it)
# if it % 200 == 0:
if it != 0 and it % 100 == 0:
fm, mae = validate(val_loader, net,
os.path.join(check_dir, 'results'),
os.path.join(opt.val_dir, 'masks'))
print(u'损失: %.4f' % (loss.item()))
print(u'最大FM: iteration %d的%.4f, 这次FM: %.4f' %
(logs['best_it'], logs['best'], fm))
logs[it] = {'FM': fm}
if fm > logs['best']:
logs['best'] = fm
logs['best_it'] = it
torch.save(net.state_dict(), '%s/net-best.pth' % (check_dir))
with open(os.path.join(check_dir, 'logs.json'), 'w') as outfile:
json.dump(logs, outfile)
torch.save(net.state_dict(), '%s/net-iter%d.pth' % (check_dir, it))
class Agent(object):
def __init__(self,
n_s,
n_a,
hiddens=(128, 64),
epsilon=1.0,
epsilon_min=0.005,
epsilon_decay=0.05,
gamma=0.99,
batch_size=64,
memory_capacity=100000,
lr=0.001,
is_dueling=False,
is_prioritize=True,
replace_iter=100,
is_soft=False,
tau=0.01,
e=0.01,
a=0.6,
b=0.4):
self.n_s = n_s
self.n_a = n_a
self.epsilon = epsilon
self.epsilon_min = epsilon_min
self.epsilon_decay = epsilon_decay
self.replace_iter = replace_iter
self.lr = lr
self.gamma = gamma
self.batch_size = batch_size
self.memory_capacity = memory_capacity
self.is_soft = is_soft
self.is_prioritize = is_prioritize
self.tau = tau
if use_gpu:
self.eval_net = Net(n_s,
n_a,
hiddens=hiddens,
is_dueling=is_dueling).cuda()
self.target_net = Net(n_s,
n_a,
hiddens=hiddens,
is_dueling=is_dueling).cuda()
else:
self.eval_net = Net(n_s,
n_a,
hiddens=hiddens,
is_dueling=is_dueling)
self.target_net = Net(n_s,
n_a,
hiddens=hiddens,
is_dueling=is_dueling)
if is_prioritize:
self.memory = Memory(memory_capacity, e, a, b)
else:
self.memory = np.zeros((memory_capacity, self.n_s * 2 + 2))
self.memory_count = 0
self.learn_count = 0
self.loss_func = nn.MSELoss()
self.optimizer = optim.Adam(self.eval_net.parameters(), lr=self.lr)
def act(self, s):
if np.random.random() <= self.epsilon:
# random
return np.random.randint(self.n_a)
else:
# max
s = FloatTensor(s)
action_value = self.eval_net(s)
a = torch.max(action_value, 1)[1].data.cpu().numpy()[0]
return a
def step(self, s, a, r, s_, done):
if self.is_prioritize:
# experience = s, a, r, s_, done
experience = np.hstack((s, [a, r], s_))
self.memory.store(experience)
self.memory_count += 1
if np.count_nonzero(self.memory.tree.tree) > self.batch_size:
tree_idx, batch, ISWeights_mb = self.memory.sample(
self.batch_size)
self.learn(batch, tree_idx, ISWeights_mb)
else:
transition = np.hstack((s, [a, r], s_))
# replace the old memory with new memory
index = self.memory_count % self.memory_capacity
self.memory[index, :] = transition
self.memory_count += 1
if self.memory_count < self.memory_capacity:
return
# sample batch transitions
sample_index = np.random.choice(self.memory_capacity,
self.batch_size)
batch = self.memory[sample_index, :]
self.learn(batch)
def learn(self, batch, tree_idx=None, ISWeights_mb=None):
b_s = torch.squeeze(FloatTensor(batch[:, :self.n_s]), 0)
b_a = torch.squeeze(LongTensor(batch[:, self.n_s:self.n_s + 1]), 0)
b_r = torch.squeeze(FloatTensor(batch[:, self.n_s + 1:self.n_s + 2]),
0)
b_s_ = torch.squeeze(FloatTensor(batch[:, -self.n_s:]), 0)
temp = self.eval_net(b_s)
eval_q = torch.gather(temp, 1, b_a)
next_max_from_eval = self.eval_net(b_s_)
next_max_from_eval_index = next_max_from_eval.max(1)[1].unsqueeze(1)
next_actions = self.target_net(b_s_).detach()
next_max = next_actions.gather(1, next_max_from_eval_index)
target_q = b_r + self.gamma * next_max # * (1 - b_done)
abs_errors = numpy(torch.sum(torch.abs(target_q - eval_q), dim=1))
loss = self.loss_func(eval_q, target_q)
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
if self.is_prioritize:
self.memory.batch_update(tree_idx=tree_idx, abs_errors=abs_errors)
self.update()
self.learn_count += 1
def update(self):
next_epsilon = self.epsilon * self.epsilon_decay
if next_epsilon <= self.epsilon_min:
self.epsilon = self.epsilon_min
else:
self.epsilon = next_epsilon
if self.is_soft:
for target_param, local_param in zip(self.target_net.parameters(),
self.eval_net.parameters()):
target_param.data.copy_(self.tau * local_param.data +
(1.0 - self.tau) * target_param.data)
else:
if self.learn_count % self.replace_iter == 0:
self.target_net.load_state_dict(self.eval_net.state_dict())
# save all net
def save(self, name):
torch.save(self.eval_net, name)
# load all net
def load(self, name):
return torch.load(name)
def main(args):
clf_model = Net().to(device)
clf_model.load_state_dict(torch.load(args.MODEL_FILE))
feature_extractor = clf_model.feature_extractor
discriminator = clf_model.classifier
critic = nn.Sequential(nn.Linear(320, 50), nn.ReLU(), nn.Linear(50, 20),
nn.ReLU(), nn.Linear(20, 1)).to(device)
half_batch = args.batch_size // 2
source_dataset = MNIST(config.DATA_DIR / 'mnist',
train=True,
download=True,
transform=Compose([GrayscaleToRgb(),
ToTensor()]))
source_loader = DataLoader(source_dataset,
batch_size=half_batch,
drop_last=True,
shuffle=True,
num_workers=0,
pin_memory=True)
target_dataset = MNISTM(train=False)
target_loader = DataLoader(target_dataset,
batch_size=half_batch,
drop_last=True,
shuffle=True,
num_workers=0,
pin_memory=True)
critic_optim = torch.optim.Adam(critic.parameters(), lr=1e-4)
clf_optim = torch.optim.Adam(clf_model.parameters(), lr=1e-4)
clf_criterion = nn.CrossEntropyLoss()
for epoch in range(1, args.epochs + 1):
batch_iterator = zip(loop_iterable(source_loader),
loop_iterable(target_loader))
total_loss = 0
total_accuracy = 0
for _ in trange(args.iterations, leave=False):
(source_x, source_y), (target_x, _) = next(batch_iterator)
# Train critic
set_requires_grad(feature_extractor, requires_grad=False)
set_requires_grad(critic, requires_grad=True)
source_x, target_x = source_x.to(device), target_x.to(device)
source_y = source_y.to(device)
with torch.no_grad():
h_s = feature_extractor(source_x).data.view(
source_x.shape[0], -1)
h_t = feature_extractor(target_x).data.view(
target_x.shape[0], -1)
for _ in range(args.k_critic):
gp = gradient_penalty(critic, h_s, h_t)
critic_s = critic(h_s)
critic_t = critic(h_t)
wasserstein_distance = critic_s.mean() - critic_t.mean()
critic_cost = -wasserstein_distance + args.gamma * gp
critic_optim.zero_grad()
critic_cost.backward()
critic_optim.step()
total_loss += critic_cost.item()
# Train classifier
set_requires_grad(feature_extractor, requires_grad=True)
set_requires_grad(critic, requires_grad=False)
for _ in range(args.k_clf):
source_features = feature_extractor(source_x).view(
source_x.shape[0], -1)
target_features = feature_extractor(target_x).view(
target_x.shape[0], -1)
source_preds = discriminator(source_features)
clf_loss = clf_criterion(source_preds, source_y)
wasserstein_distance = critic(source_features).mean() - critic(
target_features).mean()
loss = clf_loss + args.wd_clf * wasserstein_distance
clf_optim.zero_grad()
loss.backward()
clf_optim.step()
mean_loss = total_loss / (args.iterations * args.k_critic)
tqdm.write(f'EPOCH {epoch:03d}: critic_loss={mean_loss:.4f}')
torch.save(clf_model.state_dict(), 'trained_models/wdgrl.pt')
class Trainer(object):
def __init__(self, train_loader, test_loader, config):
self.config = config
self.device = config.device
self.train_loader = train_loader
self.test_loader = test_loader
self.n_epoch = config.n_epoch
self.lr = config.lr
self.gamma = config.gamma
self.device = config.device
# self.start_epoch = 1
self.start_itr = 1
n_classes = len(self.train_loader.dataset.classes)
self.model = Net(n_classes=n_classes).to(self.device)
print(self.model)
print('Initialized model...\n')
self.optim = torch.optim.Adadelta(self.model.parameters(), self.lr)
self.scheduler = StepLR(self.optim, step_size=1, gamma=self.gamma)
# if not self.config.model_state_path == '':
# self._load_models(self.config.model_state_path)
self.writer = SummaryWriter(log_dir=self.config.log_dir)
def train(self):
self.model.train()
n_itr = self.start_itr
print('Start training...!')
for epoch in range(1, self.n_epoch + 1):
with tqdm(total=len(self.train_loader)) as pbar:
for idx, (img, label) in enumerate(self.train_loader):
pbar.set_description(
f'Epoch[{epoch}/{self.n_epoch}], iteration[{idx}/{len(self.train_loader)}]'
)
img, label = img.to(self.device), label.to(self.device)
self.optim.zero_grad()
output = self.model(img)
loss = F.nll_loss(output, label)
loss.backward()
self.optim.step()
if n_itr % self.config.log_interval == 0:
pbar.set_postfix(OrderedDict(loss=loss.item()))
tqdm.write(
f'Epoch[{epoch}], iteration[{idx}/{len(self.train_loader)}], loss: {loss.item()}'
)
self.writer.add_scalar('loss/loss', loss.item(), n_itr)
if n_itr % self.config.checkpoint_interval == 0:
self._save_models(epoch, n_itr)
n_itr += 1
pbar.update()
self.scheduler.step()
self.test(n_itr)
self.writer.close()
def test(self, n_itr):
self.model.eval()
test_loss = 0
correct = 0
print('Start testing...!')
with torch.no_grad():
for _, (img, label) in enumerate(self.test_loader):
img, label = img.to(self.device), label.to(self.device)
output = self.model(img)
test_loss += F.nll_loss(output, label, reduction='sum').item()
pred = output.argmax(dim=1, keepdim=True)
correct += pred.eq(label.view_as(pred)).sum().item()
test_loss /= len(self.test_loader.dataset)
accuracy = correct / len(self.test_loader.dataset)
self.writer.add_scalar('accuracy/test_accuracy', accuracy, n_itr)
tqdm.write(
f'Test: Average loss: {test_loss}, Accuracy: {accuracy * 100.0}%')
self.model.train()
def _save_models(self, epoch, n_itr):
checkpoint_name = f'{self.config.dataset_name}_model_ckpt_{n_itr}.pt'
checkpoint_path = os.path.join(self.config.checkpoint_dir,
checkpoint_name)
torch.save(
{
# 'epoch': epoch,
'n_itr': n_itr,
'model': self.model.state_dict(),
'optim': self.optim.state_dict(),
},
checkpoint_path)
tqdm.write(f'Saved models state_dict: n_itr_{n_itr}')
def _load_models(self, model_state_path):
checkpoint = torch.load(model_state_path)
# self.start_epoch = checkpoint['epoch']
self.start_itr = checkpoint['n_itr'] + 1
self.model.load_state_dict(checkpoint['model'])
self.optim.load_state_dict(checkpoint['optim'])
print(f'start_itr: {self.start_itr}')
print('Loaded pretrained models...\n')
class dueling_agent:
def __init__(self, env, args):
# define some important
self.env = env
self.args = args
# trying to define the network
self.net = Net(self.env.action_space.n)
self.target_net = Net(self.env.action_space.n)
# make sure the target net has the same weights as the network
self.target_net.load_state_dict(self.net.state_dict())
if self.args.cuda:
self.net.cuda()
self.target_net.cuda()
# define the optimizer
self.optimizer = torch.optim.Adam(self.net.parameters(),
lr=self.args.lr)
# define the replay memory
self.buffer = replay_memory(self.args.buffer_size)
# define the linear schedule of the exploration
self.exploration_schedule = linear_schedule(int(self.args.total_timesteps * self.args.exploration_fraction), \
self.args.final_ratio, self.args.init_ratio)
# create the folder to save the models
if not os.path.exists(self.args.save_dir):
os.mkdir(self.args.save_dir)
# set the environment folder
self.model_path = os.path.join(self.args.save_dir, self.args.env_name)
if not os.path.exists(self.model_path):
os.mkdir(self.model_path)
# start to do the training
def learn(self):
episode_reward = [0.0]
obs = np.array(self.env.reset())
td_loss = 0
for timestep in range(self.args.total_timesteps):
explore_eps = self.exploration_schedule.get_value(timestep)
with torch.no_grad():
obs_tensor = self._get_tensors(obs)
action_value = self.net(obs_tensor)
# select actions
action = select_actions(action_value, explore_eps)
# excute actions
obs_, reward, done, _ = self.env.step(action)
obs_ = np.array(obs_)
# tryint to append the samples
self.buffer.add(obs, action, reward, obs_, float(done))
obs = obs_
# add the rewards
episode_reward[-1] += reward
if done:
obs = np.array(self.env.reset())
episode_reward.append(0.0)
if timestep > self.args.learning_starts and timestep % self.args.train_freq == 0:
# start to sample the samples from the replay buffer
batch_samples = self.buffer.sample(self.args.batch_size)
td_loss = self._update_network(batch_samples)
if timestep > self.args.learning_starts and timestep % self.args.target_network_update_freq == 0:
# update the target network
self.target_net.load_state_dict(self.net.state_dict())
if len(episode_reward[-101:-1]) == 0:
mean_reward_per_100 = 0
else:
mean_reward_per_100 = np.mean(episode_reward[-101:-1])
num_episode = len(episode_reward) - 1
if done and num_episode % self.args.display_interval == 0:
print('[{}] Frames: {}, Episode: {}, Mean: {:.3f}, Loss: {:.3f}'.format(datetime.now(), timestep, num_episode, \
mean_reward_per_100, td_loss))
torch.save(self.net.state_dict(),
self.model_path + '/model.pt')
# update the network
def _update_network(self, samples):
obses, actions, rewards, obses_next, dones = samples
# convert the data to tensor
obses = self._get_tensors(obses)
actions = torch.tensor(actions, dtype=torch.int64).unsqueeze(-1)
rewards = torch.tensor(rewards, dtype=torch.float32).unsqueeze(-1)
obses_next = self._get_tensors(obses_next)
dones = torch.tensor(1 - dones, dtype=torch.float32).unsqueeze(-1)
# convert into gpu
if self.args.cuda:
actions = actions.cuda()
rewards = rewards.cuda()
dones = dones.cuda()
# calculate the target value
with torch.no_grad():
q_value_temp = self.net(obses_next)
action_max_idx = torch.argmax(q_value_temp, dim=1, keepdim=True)
target_action_value = self.target_net(obses_next)
target_action_max_value = target_action_value.gather(
1, action_max_idx)
target_action_max_value = target_action_max_value.detach()
# target
expected_value = rewards + self.args.gamma * target_action_max_value * dones
# get the real q value
action_value = self.net(obses)
real_value = action_value.gather(1, actions)
loss = (expected_value - real_value).pow(2).mean()
# start to update
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
return loss.item()
# get tensors
def _get_tensors(self, obs):
if obs.ndim == 3:
obs = np.transpose(obs, (2, 0, 1))
obs = np.expand_dims(obs, 0)
elif obs.ndim == 4:
obs = np.transpose(obs, (0, 3, 1, 2))
obs = torch.tensor(obs, dtype=torch.float32)
if self.args.cuda:
obs = obs.cuda()
return obs
import torch
import torch.nn as nn
import torch.optim as optim
from models import Net
from model_train import net_sample_output, train_net
from model_evaluate import visualize_output
use_cuda = torch.cuda.is_available()
device = torch.device("cuda" if use_cuda else "cpu")
net = Net().to(device)
# define loss and optimisation functions
criterion = nn.MSELoss()
optimizer = optim.Adam(net.parameters(), lr=0.001)
# train the network
n_epochs = 100
train_net(net, device, criterion, optimizer, n_epochs)
# get a sample of test data
test_images, test_outputs, gt_pts = net_sample_output(net, device)
visualize_output(test_images, test_outputs, gt_pts)
# after training, save the model parameters in the dir 'saved_models'
torch.save(net.state_dict(), 'saved_models/keypoints_model.pt')
# zero the parameter (weight) gradients
optimizer.zero_grad()
# backward pass to calculate the weight gradients
loss.backward()
# update the weights
optimizer.step()
# print loss statistics
running_loss += loss.item()
if batch_i % 10 == 9: # print every 10 batches
print('Epoch: {}, Batch: {}, Avg. Loss: {}'.format(
epoch + 1, batch_i + 1, running_loss / 10))
running_loss = 0.0
print('Finished Training')
# train your network
n_epochs = args.epochs # start small, and increase when you've decided on your model structure and hyperparams
train_net(n_epochs)
# ## TODO: change the name to something uniqe for each new model
model_dir = 'saved_models/'
model_name = args.experiment_name + '.pt'
# # after training, save your model parameters in the dir 'saved_models'
torch.save(net.state_dict(), model_dir + model_name)
class NNet():
"""
Wrapper to manage neural net.
"""
def __init__(self, args):
self.args = args
self.num_channels = NUM_CHANNELS
if args.netType == 1:
self.net = Net(self.num_channels, args)
elif args.netType == 2:
self.net = Net2(self.num_channels, args)
if args.cuda:
self.net = self.net.cuda()
self.load_dataset_from_folder()
self.writer = SummaryWriter()
self.unique_tok = str(time.time())
self.init_weights()
def init_weights(self):
"""
Initialize by Xavier weights
"""
self.net.apply(init_weights)
def load_dataset_from_folder(self):
"""
Load complete dataset
"""
all_data_path = self.args.all_data_path
validation_split_size = self.args.validation_split_size
batch_size = self.args.batch_size
num_workers = self.args.num_workers
shuffle = self.args.shuffle
all_data = ImageFolder(root=all_data_path, transform=TRANSFORM)
classes = all_data.classes
self.classes = classes
validation_size = int(validation_split_size * len(all_data))
test_size = int(validation_split_size * len(all_data))
train_size = len(all_data) - 2 * validation_size
train_dataset, val_dataset, test_dataset = random_split(
all_data, [train_size, validation_size, test_size])
training_data_loader = DataLoader(train_dataset,
batch_size=batch_size,
num_workers=num_workers,
shuffle=shuffle)
validation_dataset_loader = DataLoader(val_dataset,
batch_size=batch_size,
num_workers=num_workers,
shuffle=shuffle)
test_dataset_loader = DataLoader(test_dataset,
batch_size=batch_size,
num_workers=num_workers,
shuffle=shuffle)
self.train_loader = training_data_loader
self.val_loader = validation_dataset_loader
self.test_loader = test_dataset_loader
def train(self):
"""
Train Neural Net
"""
if self.args.optim == 'RMSprop':
optimizer = optim.RMSprop(self.net.parameters(),
lr=self.args.lr,
momentum=self.args.momentum,
weight_decay=self.args.l2_regularization)
elif self.args.optim == 'SGD':
optimizer = optim.SGD(self.net.parameters(),
lr=self.args.lr,
momentum=self.args.momentum)
elif self.args.optim == 'Adam':
optimizer = optim.Adam(self.net.parameters(), lr=self.args.lr)
criterion = nn.CrossEntropyLoss()
# scheduler = optim.lr_scheduler.StepLR(
# optimizer, step_size=self.args.scheduler_step_size, gamma=self.args.scheduler_gamma)
self.net.train()
for epoch in range(self.args.epoch):
start_time = time.time()
running_loss_t = 0.0
num_batches = 0
y_true = []
y_pred = []
# print('Epoch: {} , LR: {}'.format(epoch+1, scheduler.get_lr()))
for data in tqdm(self.train_loader):
inputs, labels = data
labels_cp = labels.clone()
# imshow(torchvision.utils.make_grid(inputs[:,:3,:,:]))
if len(inputs) < 2:
continue
if self.args.cuda:
inputs = inputs.cuda()
labels = labels.cuda()
outputs = self.net(inputs)
loss = criterion(outputs, labels)
_, predicted = torch.max(outputs, 1)
predicted = predicted.cpu()
for i, pred in enumerate(predicted):
y_pred.append(pred)
y_true.append(labels_cp[i])
optimizer.zero_grad()
loss.backward()
optimizer.step()
running_loss_t += loss.item()
num_batches += 1
end_time = time.time()
train_f1 = f1_score(y_true, y_pred, average='weighted')
# scheduler.step()
self.save(epoch + 1)
self.writer.add_scalar('Loss/train', running_loss_t / num_batches,
epoch + 1)
self.writer.add_scalar('F1/train', train_f1, epoch + 1)
loss_v, val_f1 = self.get_validation_loss(criterion)
self.writer.add_scalar('Loss/val', loss_v, epoch + 1)
self.writer.add_scalar('F1/val', val_f1, epoch + 1)
print(
"Epoch {} Time {:.2f}s Train-Loss {:.3f} Val-Loss {:.3f} Train-F1 {:.3f} Val-F1 {:.3f}"
.format(epoch + 1, end_time - start_time,
running_loss_t / num_batches, loss_v, train_f1,
val_f1))
def get_validation_loss(self, criterion):
"""
Check validation loss
"""
running_loss = 0.0
num_batches = 0
self.net.eval()
y_true = []
y_pred = []
with torch.no_grad():
for data in tqdm(self.val_loader):
images, labels = data
labels_cp = labels.clone()
if self.args.cuda:
images = images.cuda()
labels = labels.cuda()
outputs = self.net(images)
_, predicted = torch.max(outputs, 1)
predicted = predicted.cpu()
for i, pred in enumerate(predicted):
y_pred.append(pred)
y_true.append(labels_cp[i])
loss = criterion(outputs, labels)
running_loss += loss.item()
num_batches += 1
self.net.train()
val_f1 = f1_score(y_true, y_pred, average='weighted')
return running_loss / num_batches, val_f1
def get_test_accuracy(self):
"""
Check overall accuracy of model
"""
y_true = []
y_pred = []
class_correct = list(0. for i in range(4))
class_total = list(0. for i in range(4))
with torch.no_grad():
for data in tqdm(self.test_loader):
images, labels = data
labels_cp = labels.clone()
if self.args.cuda:
images = images.cuda()
labels = labels.cuda()
outputs = self.net(images)
_, predicted = torch.max(outputs, 1)
predicted = predicted.cpu()
for i, pred in enumerate(predicted):
y_pred.append(pred)
y_true.append(labels_cp[i])
c = (predicted == labels_cp).squeeze()
for i in range(min(self.args.batch_size, len(labels_cp))):
label = labels_cp[i]
class_correct[label] += c[i].item()
class_total[label] += 1
print("Test F1: ", f1_score(y_true, y_pred, average='weighted'))
def save(self, epochs, folder_path="../models/"):
"""
Save Model
"""
dict_save = {'params': self.net.state_dict(), 'classes': self.classes}
name = folder_path + self.unique_tok + '_' + str(epochs) + '.model'
torch.save(dict_save, name)
print('Model saved at {}'.format(name))
return name
def load(self, path):
"""
Load a saved model
"""
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
dict_load = torch.load(path, map_location=torch.device(device))
self.net.load_state_dict(dict_load['params'])
return dict_load['classes']
def predict(self, inp):
"""
Predict using net
"""
if self.args.cuda:
inp = inp.cuda()
self.net.eval()
with torch.no_grad():
vals = self.net(inp)
print(vals)
_, predicted = torch.max(vals, 1)
predicted = predicted.cpu()
result_class = self.classes[predicted]
return result_class
def main(args):
model = Net().to(device)
model.load_state_dict(torch.load(args.MODEL_FILE))
feature_extractor = model.feature_extractor
clf = model.classifier
discriminator = nn.Sequential(GradientReversal(), nn.Linear(320, 50),
nn.ReLU(), nn.Linear(50, 20), nn.ReLU(),
nn.Linear(20, 1)).to(device)
half_batch = args.batch_size // 2
source_dataset = MNIST(config.DATA_DIR / 'mnist',
train=True,
download=True,
transform=Compose([GrayscaleToRgb(),
ToTensor()]))
source_loader = DataLoader(source_dataset,
batch_size=half_batch,
shuffle=True,
num_workers=1,
pin_memory=True)
target_dataset = MNISTM(train=False)
target_loader = DataLoader(target_dataset,
batch_size=half_batch,
shuffle=True,
num_workers=1,
pin_memory=True)
optim = torch.optim.Adam(
list(discriminator.parameters()) + list(model.parameters()))
for epoch in range(1, args.epochs + 1):
batches = zip(source_loader, target_loader)
n_batches = min(len(source_loader), len(target_loader))
total_domain_loss = total_label_accuracy = 0
target_label_accuracy = 0
for (source_x,
source_labels), (target_x,
target_labels) in tqdm(batches,
leave=False,
total=n_batches):
x = torch.cat([source_x, target_x])
x = x.to(device)
domain_y = torch.cat([
torch.ones(source_x.shape[0]),
torch.zeros(target_x.shape[0])
])
domain_y = domain_y.to(device)
label_y = source_labels.to(device)
features = feature_extractor(x).view(x.shape[0], -1)
domain_preds = discriminator(features).squeeze()
label_preds = clf(features[:source_x.shape[0]])
domain_loss = F.binary_cross_entropy_with_logits(
domain_preds, domain_y)
label_loss = F.cross_entropy(label_preds, label_y)
loss = domain_loss + label_loss
optim.zero_grad()
loss.backward()
optim.step()
total_domain_loss += domain_loss.item()
total_label_accuracy += (
label_preds.max(1)[1] == label_y).float().mean().item()
target_label_preds = clf(features[source_x.shape[0]:])
target_label_accuracy += (target_label_preds.cpu().max(1)[1] ==
target_labels).float().mean().item()
mean_loss = total_domain_loss / n_batches
mean_accuracy = total_label_accuracy / n_batches
target_mean_accuracy = target_label_accuracy / n_batches
tqdm.write(
f'EPOCH {epoch:03d}: domain_loss={mean_loss:.4f}, '
f'source_accuracy={mean_accuracy:.4f}, target_accuracy={target_mean_accuracy:.4f}'
)
torch.save(model.state_dict(), 'trained_models/revgrad.pt')
def main(args):
# TODO: add DTN model
model = Net().to(device)
model.load_state_dict(torch.load(args.MODEL_FILE))
feature_extractor = model.feature_extractor
clf = model.classifier
discriminator = nn.Sequential(GradientReversal(), nn.Linear(320, 50),
nn.ReLU(), nn.Linear(50, 20), nn.ReLU(),
nn.Linear(20, 1)).to(device)
half_batch = args.batch_size // 2
if args.adapt_setting == 'mnist2mnistm':
source_dataset = MNIST(config.DATA_DIR / 'mnist',
train=True,
download=True,
transform=Compose(
[GrayscaleToRgb(),
ToTensor()]))
target_dataset = MNISTM(train=False)
elif args.adapt_setting == 'svhn2mnist':
source_dataset = ImageClassdata(txt_file=args.src_list,
root_dir=args.src_root,
img_type=args.img_type,
transform=transforms.Compose([
transforms.Resize(28),
transforms.ToTensor(),
]))
target_dataset = ImageClassdata(txt_file=args.tar_list,
root_dir=args.tar_root,
img_type=args.img_type,
transform=transforms.Compose([
transforms.ToTensor(),
]))
elif args.adapt_setting == 'mnist2usps':
source_dataset = ImageClassdata(txt_file=args.src_list,
root_dir=args.src_root,
img_type=args.img_type,
transform=transforms.Compose([
transforms.ToTensor(),
]))
target_dataset = ImageClassdata(txt_file=args.tar_list,
root_dir=args.tar_root,
img_type=args.img_type,
transform=transforms.Compose([
transforms.Resize(28),
transforms.ToTensor(),
]))
else:
raise NotImplementedError
source_loader = DataLoader(source_dataset,
batch_size=half_batch,
shuffle=True,
num_workers=1,
pin_memory=True,
drop_last=True)
target_loader = DataLoader(target_dataset,
batch_size=half_batch,
shuffle=True,
num_workers=1,
pin_memory=True,
drop_last=True)
optim = torch.optim.Adam(
list(discriminator.parameters()) + list(model.parameters()))
if not os.path.exists('logs'): os.makedirs('logs')
f = open(f'logs/{args.adapt_setting}_{args.name}.txt', 'w+')
for epoch in range(1, args.epochs + 1):
batches = zip(source_loader, target_loader)
n_batches = min(len(source_loader), len(target_loader))
total_domain_loss = total_label_accuracy = 0
target_label_accuracy = 0
for (source_x,
source_labels), (target_x,
target_labels) in tqdm(batches,
leave=False,
total=n_batches):
x = torch.cat([source_x, target_x])
x = x.to(device)
domain_y = torch.cat([
torch.ones(source_x.shape[0]),
torch.zeros(target_x.shape[0])
])
domain_y = domain_y.to(device)
label_y = source_labels.to(device)
features = feature_extractor(x).view(x.shape[0], -1)
domain_preds = discriminator(features).squeeze()
label_preds = clf(features[:source_x.shape[0]])
domain_loss = F.binary_cross_entropy_with_logits(
domain_preds, domain_y)
label_loss = F.cross_entropy(label_preds, label_y)
loss = domain_loss + label_loss
optim.zero_grad()
loss.backward()
optim.step()
total_domain_loss += domain_loss.item()
total_label_accuracy += (
label_preds.max(1)[1] == label_y).float().mean().item()
target_label_preds = clf(features[source_x.shape[0]:])
target_label_accuracy += (target_label_preds.cpu().max(1)[1] ==
target_labels).float().mean().item()
mean_loss = total_domain_loss / n_batches
mean_accuracy = total_label_accuracy / n_batches
target_mean_accuracy = target_label_accuracy / n_batches
tqdm.write(
f'EPOCH {epoch:03d}: domain_loss={mean_loss:.4f}, '
f'source_accuracy={mean_accuracy:.4f}, target_accuracy={target_mean_accuracy:.4f}'
)
f.write(
f'EPOCH {epoch:03d}: domain_loss={mean_loss:.4f}, '
f'source_accuracy={mean_accuracy:.4f}, target_accuracy={target_mean_accuracy:.4f}\n'
)
torch.save(
model.state_dict(),
f'trained_models/{args.adapt_setting}_{args.name}_ep{epoch}.pt')
f.close()
class a2c_agent:
def __init__(self, envs, args):
self.envs = envs
self.args = args
# define the network
self.net = Net(self.envs.action_space.n)
if self.args.cuda:
self.net.cuda()
# define the optimizer
self.optimizer = torch.optim.RMSprop(self.net.parameters(),
lr=self.args.lr,
eps=self.args.eps,
alpha=self.args.alpha)
if not os.path.exists(self.args.save_dir):
os.mkdir(self.args.save_dir)
# check the saved path for envs..
self.model_path = self.args.save_dir + self.args.env_name + '/'
if not os.path.exists(self.model_path):
os.mkdir(self.model_path)
# get the obs..
self.batch_ob_shape = (self.args.num_processes * self.args.nsteps,
) + self.envs.observation_space.shape
self.obs = np.zeros(
(self.args.num_processes, ) + self.envs.observation_space.shape,
dtype=self.envs.observation_space.dtype.name)
self.obs[:] = self.envs.reset()
self.dones = [False for _ in range(self.args.num_processes)]
# train the network..
def learn(self):
if not self.args.no_sil:
sil_model = sil_module(self.net, self.args, self.optimizer)
num_updates = self.args.total_frames // (self.args.num_processes *
self.args.nsteps)
# get the reward to calculate other information
episode_rewards = torch.zeros([self.args.num_processes, 1])
final_rewards = torch.zeros([self.args.num_processes, 1])
# start to update
for update in range(num_updates):
mb_obs, mb_rewards, mb_actions, mb_dones = [], [], [], []
for step in range(self.args.nsteps):
with torch.no_grad():
input_tensor = self._get_tensors(self.obs)
_, pi = self.net(input_tensor)
# select actions
actions = select_actions(pi)
cpu_actions = actions.squeeze(1).cpu().numpy()
# start to store the information
mb_obs.append(np.copy(self.obs))
mb_actions.append(cpu_actions)
mb_dones.append(self.dones)
# step
obs, rewards, dones, _ = self.envs.step(cpu_actions)
# process rewards...
raw_rewards = copy.deepcopy(rewards)
rewards = np.sign(rewards)
# start to store the rewards
self.dones = dones
if not self.args.no_sil:
sil_model.step(input_tensor.detach().cpu().numpy(),
cpu_actions, raw_rewards, dones)
mb_rewards.append(rewards)
for n, done in enumerate(dones):
if done:
self.obs[n] = self.obs[n] * 0
self.obs = obs
raw_rewards = torch.from_numpy(
np.expand_dims(np.stack(raw_rewards), 1)).float()
episode_rewards += raw_rewards
# get the masks
masks = torch.FloatTensor([[0.0] if done_ else [1.0]
for done_ in dones])
final_rewards *= masks
final_rewards += (1 - masks) * episode_rewards
episode_rewards *= masks
# update the obs
mb_dones.append(self.dones)
# process the rollouts
mb_obs = np.asarray(mb_obs, dtype=np.uint8).swapaxes(1, 0).reshape(
self.batch_ob_shape)
mb_rewards = np.asarray(mb_rewards,
dtype=np.float32).swapaxes(1, 0)
mb_actions = np.asarray(mb_actions, dtype=np.int32).swapaxes(1, 0)
mb_dones = np.asarray(mb_dones, dtype=np.bool).swapaxes(1, 0)
mb_masks = mb_dones[:, :-1]
mb_dones = mb_dones[:, 1:]
with torch.no_grad():
input_tensor = self._get_tensors(self.obs)
last_values, _ = self.net(input_tensor)
# compute returns
for n, (rewards, dones, value) in enumerate(
zip(mb_rewards, mb_dones,
last_values.detach().cpu().numpy().squeeze())):
rewards = rewards.tolist()
dones = dones.tolist()
if dones[-1] == 0:
rewards = discount_with_dones(rewards + [value],
dones + [0],
self.args.gamma)[:-1]
else:
rewards = discount_with_dones(rewards, dones,
self.args.gamma)
mb_rewards[n] = rewards
mb_rewards = mb_rewards.flatten()
mb_actions = mb_actions.flatten()
# start to update network
vl, al, ent = self._update_network(mb_obs, mb_rewards, mb_actions)
# start to update the sil_module
if not self.args.no_sil:
mean_adv, num_samples = sil_model.train_sil_model()
if update % self.args.log_interval == 0:
if not self.args.no_sil:
print('[{}] Update: {}/{}, Frames: {}, Rewards: {:.2f}, VL: {:.3f}, PL: {:.3f},' \
'Ent: {:.2f}, Min: {}, Max:{}, BR:{}, E:{}, VS:{}, S:{}'.format(\
datetime.now(), update, num_updates, (update+1)*(self.args.num_processes * self.args.nsteps),\
final_rewards.mean(), vl, al, ent, final_rewards.min(), final_rewards.max(), sil_model.get_best_reward(), \
sil_model.num_episodes(), num_samples, sil_model.num_steps()))
else:
print('[{}] Update: {}/{}, Frames: {}, Rewards: {:.2f}, VL: {:.3f}, PL: {:.3f},' \
'Ent: {:.2f}, Min: {}, Max:{}'.format(\
datetime.now(), update, num_updates, (update+1)*(self.args.num_processes * self.args.nsteps),\
final_rewards.mean(), vl, al, ent, final_rewards.min(), final_rewards.max()))
torch.save(self.net.state_dict(), self.model_path + 'model.pt')
# update_network
def _update_network(self, obs, returns, actions):
# evaluate the actions
input_tensor = self._get_tensors(obs)
values, pi = self.net(input_tensor)
# define the tensor of actions, returns
returns = torch.tensor(returns, dtype=torch.float32).unsqueeze(1)
actions = torch.tensor(actions, dtype=torch.int64).unsqueeze(1)
if self.args.cuda:
returns = returns.cuda()
actions = actions.cuda()
# evaluate actions
action_log_probs, dist_entropy = evaluate_actions(pi, actions)
# calculate advantages...
advantages = returns - values
# get the value loss
value_loss = advantages.pow(2).mean()
# get the action loss
action_loss = -(advantages.detach() * action_log_probs).mean()
# total loss
total_loss = action_loss + self.args.value_loss_coef * value_loss - self.args.entropy_coef * dist_entropy
# start to update
self.optimizer.zero_grad()
total_loss.backward()
torch.nn.utils.clip_grad_norm_(self.net.parameters(),
self.args.max_grad_norm)
self.optimizer.step()
return value_loss.item(), action_loss.item(), dist_entropy.item()
# get the tensors...
def _get_tensors(self, obs):
input_tensor = torch.tensor(np.transpose(obs, (0, 3, 1, 2)),
dtype=torch.float32)
if self.args.cuda:
input_tensor = input_tensor.cuda()
return input_tensor
def train():
# load dataset
train_loader, test_loader = load_data()
channels = 1
class_num = 10
# params
lr = 0.01
momentum = 0.5
model = Net(channels, class_num).to(device)
#optimizer = optim.Adam(model.parameters(), lr=lr)
optimizer = optim.SGD(model.parameters(), lr=lr, momentum=momentum)
epochs = 30
# train data
model.train()
all_train_correct = []
all_test_correct = []
all_train_loss = []
all_test_loss = []
for epoch in range(epochs):
train_correct = 0
train_loss = 0
print(len(train_loader))
for batch_idx, (data, target) in enumerate(train_loader):
data, target = data.to(device), target.to(device)
optimizer.zero_grad() # zero to gradient
output = model(data)
loss = F.cross_entropy(output, target) # loss
loss.backward()
optimizer.step()
pred = output.argmax(dim=1, keepdim=True)
train_correct += pred.eq(target.view_as(pred)).sum().item()
train_loss += loss.item()
if batch_idx % 50 == 0:
print("train Epochs %d %d/%d loss %.6f" %
(epoch, batch_idx, len(train_loader), loss.item()))
test_loss = 0
test_correct = 0
with torch.no_grad():
for batch_idx, (data, target) in enumerate(test_loader):
data, target = data.to(device), target.to(device)
output = model(data)
test_loss += F.cross_entropy(output, target)
pred = output.argmax(dim=1, keepdim=True)
test_correct += pred.eq(target.view_as(pred)).sum().item()
print("Test average loss %.6f test correct %.6f train correct %.6f" %
(test_loss / len(test_loader), test_correct * 1.0 /
(len(test_loader) * test_batch_size), train_correct * 1.0 /
(len(train_loader) * batch_size)))
torch.save(
model.state_dict(), "model/model_cnn_epoch_" + str(epoch) + '_' +
str(test_correct * 1.0 /
(len(test_loader) * test_batch_size)) + ".pt")
all_train_correct.append(train_correct * 1.0 /
(len(train_loader) * batch_size))
all_test_correct.append(test_correct * 1.0 /
(len(test_loader) * test_batch_size))
all_train_loss.append(train_loss)
all_test_loss.append(test_loss)
print('all_train_correct', all_train_correct)
print('all_test_correct', all_test_correct)
print('all_train_loss', all_train_loss)
print('all_test_loss', all_test_loss)
